如何加速Numba实现的get_cum_dist函数？其C++-O3版本快10倍

原Python函数定义 #

from numba import njit
import numpy as np

@njit(fastmath=True)
def get_cum_dist(perm: np.ndarray, c: np.ndarray, n: int) -> np.array:
    cum_dist = np.empty(n)
    cum_dist[0] = 0.
    cum_dist[1] = 0.
    for i in range(1, n - 1):
        cum_dist[i + 1] = cum_dist[i] + c[perm[i - 1], perm[i]]
    return cum_dist

输入示例 #

n = 1000
perm = np.random.permutation(n)
c = np.random.random((n+1,n+1))
cum_dist = get_cum_dist(perm, c, n)

该函数在算法中被频繁调用，需要可行的性能优化方案。我自己尝试了C++实现（未对输入做随机置换），开启-O3编译优化后，运行速度比Numba版本快10倍，代码如下：

测试用C++实现 #

#include <iostream>
using namespace std;
#include <bits/stdc++.h>
#include <chrono>
using namespace std::chrono;

int main() {
    int n, sum = 0;
    int rows = 1001;
    int cols = 1001;
    int randArr[rows][cols];
    for (int i=0;i<rows;i++)
        for (int j=0; j<cols; j++)
            randArr[i][j] = 1 + (rand() % 500);
    n = 1000;
    int arr[1000]={0};
    auto start = high_resolution_clock::now();
    for (int i = 1; i <= n; ++i) {
        arr[i+1]=arr[i] + randArr[i-1][i];
    }
    auto stop = high_resolution_clock::now();
    auto duration = duration_cast<nanoseconds>(stop - start);
    cout << duration.count() << endl;
    return 0;
}

具体优化建议 #

针对Numba版本的优化 #

1. 调整内存访问模式
   原代码中c的二维随机索引是性能瓶颈，会导致非连续内存访问、缓存命中率低。可以预先提取需要的c元素为一维数组，再进行累加：

   @njit(fastmath=True)
   def get_cum_dist_opt(perm: np.ndarray, c_flat: np.ndarray, n: int) -> np.array:
       cum_dist = np.empty(n)
       cum_dist[0] = 0.
       cum_dist[1] = 0.
       for i in range(1, n-1):
           cum_dist[i+1] = cum_dist[i] + c_flat[i-1]
       return cum_dist
   # 调用前预处理
   c_flat = c[perm[:-1], perm[1:]]
   cum_dist = get_cum_dist_opt(perm, c_flat, n)
   

2. 指定数组类型与内存对齐
   显式指定数组类型（比如perm用np.int32、c用np.float32，精度允许时），减少类型转换开销；用np.empty(n, dtype=np.float32)预分配对齐内存，提升访问效率。

3. 精简循环逻辑
   移除冗余参数n，直接通过perm.shape[0]获取；同时利用前两个元素为0的初始化值，从i=2开始循环，减少循环次数。

跨语言/替代方案优化 #

1. 完善C++实现并封装为Python扩展
   补全C++代码的随机置换逻辑，用pybind11或Cython封装为Python可调用函数。编译时添加-O3 -march=native选项，最大化编译器优化。核心逻辑示例：

   void get_cum_dist_cpp(const std::vector<int>& perm, const std::vector<std::vector<double>>& c, std::vector<double>& cum_dist) {
       int n = perm.size();
       cum_dist.resize(n);
       cum_dist[0] = 0.0;
       cum_dist[1] = 0.0;
       for (int i = 1; i < n-1; ++i) {
           cum_dist[i+1] = cum_dist[i] + c[perm[i-1]][perm[i]];
       }
   }
   

2. 使用NumPy向量化操作
   直接用np.cumsum替代循环，无需编译且性能接近Numba版本：

   def get_cum_dist_np(perm, c, n):
       cum_dist = np.zeros(n)
       c_flat = c[perm[:-1], perm[1:]]
       cum_dist[2:] = np.cumsum(c_flat[:n-2])
       return cum_dist
   

3. 切换至PyPy解释器
   PyPy对纯Python循环的优化优于CPython，若项目允许切换解释器，配合Numba可进一步提升性能（注意检查Numba与PyPy的兼容性）。

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