从Numba获取生成的LLVM

Question

在用Numba编译一个Python函数之后，比如：

from numba import jit

@jit
def sum(x, y):
    return x + y

我该如何获取这个编译函数生成的LLVM代码（以字符串形式）呢？

看起来在Numba的早期版本中，可以通过编译函数的lfunc属性来获取这个信息，但现在似乎不再有效了。

类似的功能之前也可以通过导出生成的LLVM汇编代码（在编译过程中）来实现。不过现在这也似乎不太管用了——除非我做错了什么。其实我不太想通过终端命令来运行，因为我更希望在Python内部获取这些代码，虽然我知道可以通过子进程来实现。

我希望能在运行时创建一个可移植的Python代码版本，这样可以进行翻译；如果有相关的建议，我非常欢迎。

谢谢

Answer 1

为了记录，从 numba 版本 0.18.0 开始（这是相关链接），获取 LLVM IR 和汇编代码的标准方法是通过在编译后的函数上调用 inspect_llvm 和 inspect_asm，比如说

@jit(nopython=True,nogil=True)
def mysum(a,b):
     return a+b

# First run the function with arguments for the code to get generated
a, b = np.random.rand(10), np.random.rand(10)

# Get the llvm IR
mysum.inspect_llvm()
# Get the assembly code
mysum.inspect_asm()

需要注意的是，返回的值是以字典的格式呈现的。要以可读的（源代码）格式显示它，只需这样做

for v, k in mysum.inspect_llvm().items():
    print(v, k)

这将输出一段很长的代码

define i32 @__main__9mysum_...
    ...
    ...
    ...

有趣的是，作为一个附带说明，生成的汇编代码显示在上面的函数中，for 循环被 LLVM 完全展开了

vmovsd  (%r10), %xmm0
vaddsd  (%rcx), %xmm0, %xmm0
vmovsd  %xmm0, -32(%rbx,%rbp)
vmovsd  (%r10), %xmm0
vaddsd  (%rcx), %xmm0, %xmm0
vmovsd  %xmm0, -24(%rbx,%rbp)
vmovsd  (%r10), %xmm0
vaddsd  (%rcx), %xmm0, %xmm0
vmovsd  %xmm0, -16(%rbx,%rbp)
vmovsd  (%r10), %xmm0
vaddsd  (%rcx), %xmm0, %xmm0
vmovsd  %xmm0, -8(%rbx,%rbp)
vmovsd  (%r10), %xmm0
vaddsd  (%rcx), %xmm0, %xmm0
vmovsd  %xmm0, (%rbx,%rbp)
vmovsd  (%r10), %xmm0
vaddsd  (%rcx), %xmm0, %xmm0
vmovsd  %xmm0, 8(%rbx,%rbp)
vmovsd  (%r10), %xmm0
vaddsd  (%rcx), %xmm0, %xmm0
vmovsd  %xmm0, 16(%rbx,%rbp)
vmovsd  (%r10), %xmm0
vaddsd  (%rcx), %xmm0, %xmm0
vmovsd  %xmm0, 24(%rbx,%rbp)
addq    $8, %rdx
addq    $64, %rbx
cmpq    %rdx, %rdi
jne LBB0_48    

Answer 2

我不太记得这是不是最好的方法，但如果你不介意以不同的方式编译这个函数，可以试试下面的方法：

from numba.compiler import compile_isolated

# second argument specifies the argument types to the sum function
cfunc = compile_isolated(sum, (types.int64, types.int64))

# get llvm IR as string
llvm_code_str = str(cfunc.llvm_module)

需要指定参数类型，因为这个函数实际上在知道它的签名之前是不会被编译的（可以通过明确指定或者实际调用这个函数来实现）。

你仍然可以通过设置环境变量NUMBA_DUMP_LLVM=1，然后运行你的Python脚本，从终端导出llvm IR（或者使用安装在Anaconda的bin路径下的numba命令，或者在Numba库的bin目录下：numba --dump-llvm test.py）。