[mojo-compiler] CompTime interpreter should be able to fold `pop.call_llvm_intrinsic`

[soraros]

Bug description

math.bit functions doesn't run at compile time.

Steps to reproduce

Consider the following code:

import math.bit as bit

fn main():
  alias n = bit.ctlz(10)

It produces the following error message:

/__w/modular/modular/Kernels/mojo/builtin/_startup.mojo:70:1: error: no viable expansions found
/__w/modular/modular/Kernels/mojo/builtin/_startup.mojo:70:1: note:   call expansion failed - no concrete specializations
/__w/modular/modular/Kernels/mojo/builtin/_startup.mojo:51:1: note:     no viable expansions found
/__w/modular/modular/Kernels/mojo/builtin/_startup.mojo:65:57: note:       call expansion failed - no concrete specializations
/__w/modular/modular/Kernels/mojo/builtin/_startup.mojo:29:1: note:         no viable expansions found
/__w/modular/modular/Kernels/mojo/builtin/_startup.mojo:42:18: note:           call expansion failed - no concrete specializations
/__w/modular/modular/Kernels/mojo/builtin/_startup.mojo:62:5: note:             no viable expansions found
/__w/modular/modular/Kernels/mojo/builtin/_startup.mojo:63:22: note:               call expansion failed - no concrete specializations
/workspaces/ubuntu/hello.mojo:3:1: note:                 no viable expansions found
def main():
^
/workspaces/ubuntu/hello.mojo:4:3: note:                   failed to evaluate 'apply'
  alias n = bit.ctlz(10)
  ^
/__w/modular/modular/Kernels/mojo/math/bit.mojo:81:1: note:                     failed to interpret function @$math::$bit::ctlz($builtin::$int::Int)_concrete
/__w/modular/modular/Kernels/mojo/math/bit.mojo:90:44: note:                       failed to fold operation pop.call_llvm_intrinsic(10 : index, #pop<simd false> : !pop.scalar<bool>)
mojo: error: failed to run the pass manager

System information

- OS: Docker on Intel macOS
- mojo -v: mojo 0.3.0 (f64f9601)
- modular -v: modular 0.1.4 (6b54d308)

[Mogball]

Thanks for filing! We haven't implemented compile-time interpretation of LLVM intrinsics, and it's not clear that's a worthwhile endeavor. That said, there are other ways we can make this work.

[soraros]

@Mogball Regardless of how the internal works, I'd still expect this code to run. Because from a users perspective, math.bit just look like any other module. If they definitely can't run at compile time, we should somehow use the type system or some decorators to communicate that fact. It would still be great if you could try "the other way" to make this work though.

I was trying to implement power_of_2 with the math.bit functions in constrained[power_of_2(n)]() for SIMD width checking, which I think is pretty valid use case.

[Mogball]

A decorator would be useful from a library design perspective to clearly communicate to users what functions are not intended to run at compile time. However, the language can't rely on library authors always being meticulous with using the decorator, which is why it provides a detailed, if cryptic, error message.

I do agree that the category of functions in math.bit should run at compile time though.

[abduld]

FYI: There is an is_power_of_2 function in the math module which should be compile-time evaluated . You can implement it yourself using

fn is_power_of_2(val: Int) -> Bool:
    return (val & (val - 1) == 0) and (val != 0)

[soraros]

@abduld Thanks for the pointer. I'm familiar with this implementation yet didn't know the existence of this function in the stdlib.

[JoeLoser]

I've renamed this issue to track the general problem from my observations this past week. From our internal Mojo discussion on this yesterday, we do want to implement support for this (folding llvm_intrinsincs so they work in compile time).

Some notes I'll share from our internal discussion:

1. We don’t want a split source of truth that causes divergence between comptime and runtime. ⇒ We don’t want a “if running in the comptime interpreter” parameter check.
2. Not all llvm intrinsics have folders and some have side effects, so not everything can be run at comptime. This is also true of “printf” or c function at comptime today.
3. That said, some simple things like this do, and can be folded, so we just need to wire that up.
4. MLIR LLVM dialect may/probably doesn’t have the ability to do this, so we might have to implement these folders upstream.

The big issue is probably plumbing (4) correctly from the list above.

[Mogball]

We can inject fold/interpret hooks into the LLVMDialect. We're not bound by upstream in this regard. I suspect there will be resistance to adding intrinsics folders upstream anyways.

[msaelices]

@JoeLoser, somehow, above's code seems to be working now. I've tested it, and it worked:

from bit import count_leading_zeros

fn main():
  alias n = count_leading_zeros(10)

Note: this is the code adapted from the description based on some bit package changes, as now it's not part of the math package and the ctlz is now called count_leading_zeros

However, I'm not sure if the bug is fixed.

[soraros]

@msaelices I think it's because it's being eliminated as dead code. This still fails:

fn main():
  alias n = count_leading_zeros(10)
  print(n)