Seeing the MLIR/LLVM or ASM code generated by Mojo

Mojo
5

You can use the compile module as follows:

import compile

fn f(x: Int) -> Int:
  return x + 1

fn main():
  print(compile._internal_compile_code[f, emission_kind="asm"]())
  print(compile._internal_compile_code[f, emission_kind="llvm"]())

However, writing code like this manually quickly becomes tedious. To make it easier, here is my little utility script to dump the IR to a file. Below is an example demonstrating how to use this lib:

import math
from sys.intrinsics import assume

from ir_utils import dump_ir

fn main():
  dump_ir[f, "out1"]()
  dump_ir[g, "out2"]()

@export  # use `export` so get cleaner names
fn f(x: Int) -> Int:
  assume(0 <= x < 100)
  return max(1, x * 2)

@export
fn g(x: Int) -> Int:
  assume(0 <= x < 100)
  return x * 2 + int(x == 0)
6 Likes