如何在Numba即时编译代码中包装NumPy函数并实现跨会话持久化磁盘缓存？

完全懂你的痛点！用objmode包装NumPy原生函数虽然能用上它的极致优化，但每次重启内核后缓存就失效，之前的编译成果全白费——这确实够闹心的。问题根源就在于Numba的缓存机制会把Python函数指针作为缓存键的一部分，而像np.sort这类函数的指针在不同会话里会变化，导致缓存命中失败，只能重新编译。

下面给你两个靠谱的解决方案，都是绕开Python层面的函数指针，直接绑定到NumPy底层稳定的C级符号，这样跨会话也能正常命中磁盘缓存：

方法一：通过ctypes获取NumPy底层C函数的稳定符号 #

NumPy的核心实现都是C写的，我们可以直接拿到它的C级函数符号——这个符号的标识在不同会话里是稳定的（只要你不换NumPy版本），Numba缓存就能基于这个稳定标识命中之前的编译结果。

具体代码示例如下：

import numpy as np
from numba import njit
from numba.extending import get_cython_function_address
import ctypes

# 获取NumPy底层排序函数的C级地址
sort_c_addr = get_cython_function_address("numpy.core._multiarray_umath", "PyArray_Sort")

# 定义该C函数的调用签名（对应PyArray_Sort的参数要求）
sort_c_func = ctypes.CFUNCTYPE(
    None,  # 返回值为void
    ctypes.c_void_p,  # 输入数组的内存指针
    ctypes.c_int,     # 排序的轴（-1表示最后一维）
    ctypes.c_int,     # 排序顺序：0为升序，1为降序
    ctypes.c_void_p   # 临时工作区，传NULL即可
)(sort_c_addr)

@njit(cache=True)
def sort_numpy_persistent(arr: np.ndarray) -> np.ndarray:
    # 先复制输入数组，避免修改原数据
    out_arr = arr.copy()
    # 调用底层C排序函数
    sort_c_func(out_arr.ctypes.data, -1, 0, None)
    return out_arr

这个方法里，我们直接绑定到NumPy底层的PyArray_Sort函数符号，它的标识不会随Python会话变化，所以Numba的磁盘缓存会被正确识别，重启内核后也不用重新编译。

方法二：用Numba Intrinsic直接调用C级函数 #

如果你想更深入Numba的编译流程，可以用Numba的intrinsic扩展，直接在LLVM编译层面调用C级函数，同样基于稳定的C符号，彻底避开Python函数指针的问题。

代码示例如下：

import numpy as np
from numba import njit, types
from numba.core import cgutils
from numba.extending import intrinsic

@intrinsic
def numpy_sort_impl(typingctx, arr):
    # 定义函数的类型签名：输入数组，返回同类型数组
    sig = arr.copy()(arr)
    
    def codegen(context, builder, sig, args):
        arr_val, = args
        # 从Numba的数组对象中提取内存指针
        data_ptr = builder.extract_value(arr_val, [0])
        # 获取PyArray_Sort的C函数地址
        sort_c_addr = context.get_function_address(
            "numpy.core._multiarray_umath", "PyArray_Sort"
        )
        # 定义C函数的类型
        sort_ftype = context.get_function_type(
            types.void, [types.voidptr, types.int32, types.int32, types.voidptr]
        )
        # 获取可调用的函数对象
        sort_c_func = builder.get_callable(sort_c_addr, sort_ftype)
        # 调用排序函数：轴为-1，升序，无临时工作区
        builder.call(
            sort_c_func,
            [
                data_ptr,
                cgutils.int32_t(typingctx, -1),
                cgutils.int32_t(typingctx, 0),
                cgutils.nullptr(typingctx)
            ]
        )
        return arr_val
    
    return sig, codegen

@njit(cache=True)
def sort_numpy_persistent(arr: np.ndarray) -> np.ndarray:
    out_arr = arr.copy()
    numpy_sort_impl(out_arr)
    return out_arr

这种方式完全在Numba的编译流程内处理，没有引入ctypes的额外开销，同样能保证跨会话的缓存有效性。

注意事项 #

• 这两种方法都依赖NumPy的底层C API，所以必须保证NumPy版本不变——如果升级NumPy，底层函数符号可能会变化，需要重新验证。
• 示例中针对的是一维数组的排序，如果你需要处理多维数组，可以调整PyArray_Sort的轴参数（比如传入0表示按第0轴排序）。
• 只要Numba和NumPy版本稳定，跨会话的磁盘缓存就能正常工作，不用再每次重启都重新编译。

备注：内容来源于stack exchange，提问作者Olibarer