Numba：在njit编译函数中调用列表内jit编译函数报错求助

环境：

• Python 3.11.4
• numba 0.57.1

我正在构建一个“编译器”类，用于从现有代码生成numba njit编译函数，代码如下：

import numba
from typing import Callable

class Foo:

    def __init__(self) -> None:
        
        self.compiled: dict[str, Callable] = dict()

        self.compile_bar()

    
    def compile_bar(self):

        @numba.njit
        def func1(a, b):
            return a * b

        @numba.njit
        def func2(a, b):
            return a / b

        l_c: list[Callable] = [func1, func2]

        @numba.njit
        def wrapper(a, b):
            for func in l_c:
                print(func(a, b))

        self.compiled.update({'bar': wrapper})

F = Foo()
print(F.compiled['bar'](1.,2.))

执行最后一行时抛出错误：

Exception has occurred: NumbaNotImplementedError X
Failed in nopython mode pipeline (step: native lowering)
<numba.core.base.OverloadSelector object at 0x7f996cd00c50>, (List(type(CPUDispatcher(<function Foo.compile_bar.<locals>.func1 at 0x7f996bcbcea0>)), True),)
During: lowering "$6load_deref.0 = freevar(l_c: [CPUDispatcher(<function Foo.compile_bar.<locals>.func1 at 0x7f996bcbcea0>), CPUDispatcher(<function Foo.compile_bar.<locals>.func2 at 0x7f99264c2de0>)])" at /workdir/porepy/scripts_vl/playground.py (33)
  File "/workdir/scripts/playground.py", line 40, in <module>
    print(F.compiled['bar'](1.,2.))
          ^^^^^^^^^^^^^^^^^^^^^^^^
numba.core.errors.NumbaNotImplementedError: Failed in nopython mode pipeline (step: native lowering)
<numba.core.base.OverloadSelector object at 0x7f996cd00c50>, (List(type(CPUDispatcher(<function Foo.compile_bar.<locals>.func1 at 0x7f996bcbcea0>)), True),)
During: lowering "$6load_deref.0 = freevar(l_c: [CPUDispatcher(<function Foo.compile_bar.<locals>.func1 at 0x7f996bcbcea0>), CPUDispatcher(<function Foo.compile_bar.<locals>.func2 at 0x7f99264c2de0>)])" at /workdir/scripts/playground.py (33)

注：func1和func2仅为简化示例，实际是基于Foo类输入参数创建的njit编译函数列表。我尝试使用numba typed list：

l_c = numba.typed.List.empty_list(numba.float64(numba.float64, numba.float64).as_type())
l_c.append(func1)
l_c.append(func2)

但结果并未改变。我认为问题出在通过列表索引访问函数，直接访问非列表中的函数可正常运行，报错信息较为晦涩，恳请提供解决方案及原因解释。

Numba的nopython模式不支持直接在编译函数中引用包含已编译函数（CPUDispatcher对象）的自由变量列表。即使使用typed list，因为列表元素是CPUDispatcher实例，Numba无法在编译期确定这些函数的类型签名并完成native lowering流程，因此抛出NumbaNotImplementedError。

方案1：使用编译期循环展开（适用于固定函数列表场景） #

利用numba.literal_unroll在编译期展开循环，让Numba直接处理每个函数，避免动态遍历函数列表：

import numba
from typing import Callable
from numba import literal_unroll

class Foo:

    def __init__(self) -> None:
        self.compiled: dict[str, Callable] = dict()
        self.compile_bar()

    def compile_bar(self):
        @numba.njit
        def func1(a, b):
            return a * b

        @numba.njit
        def func2(a, b):
            return a / b

        l_c = [func1, func2]

        @numba.njit
        def wrapper(a, b):
            # 编译期展开循环，直接处理每个函数
            for func in literal_unroll(l_c):
                print(func(a, b))

        self.compiled.update({'bar': wrapper})

F = Foo()
print(F.compiled['bar'](1.,2.))

方案2：使用函数指针数组（适用于动态生成的函数列表） #

将已编译函数转为Numba支持的函数指针，存入typed列表后调用：

import numba
from typing import Callable
from numba.types import float64
from numba import njit

# 定义统一的函数类型：输入两个float64，返回float64
func_type = float64(float64, float64)

class Foo:

    def __init__(self) -> None:
        self.compiled: dict[str, Callable] = dict()
        self.compile_bar()

    def compile_bar(self):
        @njit(func_type)
        def func1(a, b):
            return a * b

        @njit(func_type)
        def func2(a, b):
            return a / b

        # 创建存储函数指针的typed列表
        func_ptrs = numba.typed.List.empty_list(func_type)
        func_ptrs.append(func1.address)
        func_ptrs.append(func2.address)

        @njit
        def wrapper(a, b):
            for ptr in func_ptrs:
                # 通过函数指针调用目标函数
                func = numba.cfunc(func_type).from_address(ptr)
                print(func(a, b))

        self.compiled.update({'bar': wrapper})

F = Foo()
print(F.compiled['bar'](1.,2.))

方案3：动态生成wrapper代码（最灵活，支持高度动态的函数列表） #

通过字符串拼接生成wrapper代码，动态编译后使用：

import numba
from typing import Callable

class Foo:

    def __init__(self) -> None:
        self.compiled: dict[str, Callable] = dict()
        self.compile_bar()

    def compile_bar(self):
        @numba.njit
        def func1(a, b):
            return a * b

        @numba.njit
        def func2(a, b):
            return a / b

        funcs = [func1, func2]
        # 拼接wrapper代码字符串
        wrapper_code = """
@numba.njit
def wrapper(a, b):
"""
        for func in funcs:
            wrapper_code += f"    print({func.__name__}(a, b))\n"

        # 执行代码并获取编译后的wrapper函数
        local_vars = {'numba': numba, 'func1': func1, 'func2': func2}
        exec(wrapper_code, globals(), local_vars)
        wrapper = local_vars['wrapper']

        self.compiled.update({'bar': wrapper})

F = Foo()
print(F.compiled['bar'](1.,2.))

内容的提问来源于stack exchange，提问作者hajdukv