如何避免Numba objmode返回非连续数组引发性能警告？

问题 #

在Numba的nopython模式中使用objmode时，我收到了NumbaPerformanceWarning，提示使用了非连续数组，原因似乎是objmode返回的是对象视图。

以下是复现问题的示例代码：

import numpy as np
from numba import njit, objmode
from scipy.linalg import expm

@njit
def foo(A, b):
    # 创建连续数组
    G_t = np.zeros(shape=A.shape, dtype=np.float64)

    with objmode(G_t="float64[:, :]"):
        G_t += expm(A)
        # 取消注释后仍存在NumbaPerformanceWarning
        # G_t = np.ascontiguousarray(G_t)

    # 取消注释后不再有NumbaPerformanceWarning
    # G_t = np.ascontiguousarray(G_t)
    return G_t @ b

N = 4
A = np.random.random(size=(N, N))
b = np.random.random(size=(N))
foo(A, b)
# NumbaPerformanceWarning: '@' is faster on contiguous arrays, called on (array(float64, 2d, A), readonly array(float64, 1d, C))
# return G_t @ b

疑问 #

我的objmode使用方式是否有误？有没有简洁的方法来避免这个性能警告？

解答 #

你的objmode用法本身没有错误，但要注意：objmode块是临时切换到Python解释器执行代码，返回的数组回到nopython模式时，Numba无法自动识别其内存连续性——哪怕你在块内调用np.ascontiguousarray，也会因为模式切换的变量传递机制，让Numba丢失连续性标记。

有两种简洁的解决方式：

• 方式一：在objmode块外部转换为连续数组
  直接启用你代码中注释的那行，把G_t = np.ascontiguousarray(G_t)放在objmode块结束后、矩阵乘法之前。Numba能正确识别转换后的连续数组，直接消除警告。
• 方式二：在objmode内直接生成连续数组
  调整objmode内的逻辑，跳过原数组的加法操作，直接生成并返回连续数组：

  with objmode(G_t="float64[:, :]"):
      G_t = np.ascontiguousarray(expm(A))
  

  这样objmode返回的本身就是连续数组，后续矩阵乘法不会触发警告。

补充说明：Numba的nopython模式依赖对内存布局的静态推断，而objmode的Python执行环境打破了这个推断链，所以必须通过显式转换来告知Numba数组的连续状态，才能让矩阵乘法这类依赖连续内存的操作发挥最优性能。

内容的提问来源于stack exchange，提问作者Louis-Amand