Seeing the MLIR/LLVM or ASM code generated by Mojo

[Dasor]

Hi,

I was wondering if there is any way to see the IR mojo generates or at lest the ASM code it generates like when you use clang -S or clang -S -emit-llvm. I have checked the help of mojo build but I haven’t seen anything.

I have some code that is giving me huge performance differences based on a parameter and I want to see what code is being generated for each parameter value.

Any help on the topic is highly appreciated.

[Caroline]

Hey @Dasor, thanks for sharing this question! @joe is out today but I’ll ask him to help with this question when he’s back tomorrow

[clattner]

This is something that Mojo is designed to do, but we haven’t gotten around to polishing it and have been perhaps “overly bashful” about making it public. That said, it is available in an undocumented compile module, in a way that could at least be interesting to poke around with to get llvm ir etc.

@sora I think you’ve been experimenting with this, could you share an example or two to show how this works?

-Chris

[sora]

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)

[clattner]

Awesome, thank you Sora!