Cython中基于外部索引数组遍历矩阵元素的性能劣于直接遍历的原因及优化方案

你的测试结果里run1()比run2()慢接近一倍，核心原因在于内存访问模式的额外开销和循环变量的来源差异：

1. 额外的内存读取操作
   在run1()中，每次循环都需要从self.indices数组里读取ix和iy两个整数，才能去索引a和b的元素。而run2()里的ix和iy是循环计数器，直接存储在CPU寄存器中，完全不需要从内存加载——这节省了两次内存读取的开销，减少了内存带宽的占用。

2. 缓存命中率的间接影响
   虽然你的indices数组是按行优先顺序生成的（和run2()的遍历顺序一致），但对indices的访问会占用CPU缓存空间，挤占a和b数组的缓存位置，导致缓存命中率略有下降。尤其是当数组规模变大时，这种影响会更明显。

3. 索引计算的额外步骤
   即使indices的遍历顺序和run2()一致，run1()中用ix和iy去索引二维数组时，Cython需要计算ix * self.a.shape[1] + iy来得到实际的内存偏移量；而run2()中，内层循环的iy是连续递增的，编译器可以优化这个偏移量计算（比如直接用自增的偏移量，不需要每次乘法）。

要在使用外部索引结构的前提下接近run2()的性能，我们需要优化内存访问模式，减少额外开销。以下是几种可行的方案：

方案1：使用扁平化的一维索引 #

将二维索引转换成扁平化的一维索引（利用numpy数组的C顺序存储特性，flat_idx = ix * ny + iy），这样可以将两次内存读取（ix和iy）减少为一次，同时简化索引计算。

修改后的Cython代码 #

%%cython -f -c=-O2 -I./
import numpy as np
cimport numpy as np
cimport cython

cdef class Test:
    cdef double[:, ::1] a, b
    cdef Py_ssize_t[::1] flat_indices  # 一维扁平化索引
    cdef Py_ssize_t ny  # 提前存储a的列数，避免循环中重复读取

    def __cinit__(self, a, b, flat_indices):
        self.a = a
        self.b = b
        self.flat_indices = flat_indices
        self.ny = a.shape[1]

    @cython.boundscheck(False)
    @cython.nonecheck(False)
    @cython.wraparound(False)
    @cython.initializedcheck(False)
    cpdef void run_flat(self):
        """使用扁平化索引"""
        cdef Py_ssize_t idx, flat_idx, ix, iy
        cdef int n = self.flat_indices.shape[0]
        cdef double* a_ptr = &self.a[0, 0]  # 直接用指针访问，更快
        cdef double* b_ptr = &self.b[0, 0]

        for idx in range(n):
            flat_idx = self.flat_indices[idx]
            # 从flat_idx反推ix和iy（如果需要用到它们的话）
            ix = flat_idx // self.ny
            iy = flat_idx % self.ny
            b_ptr[flat_idx] = ix * iy * a_ptr[flat_idx]

Python侧调用代码 #

import itertools
import numpy as np
N = 256
a = np.random.rand(N, N)
b = np.zeros_like(a)
# 生成扁平化索引
flat_indices = np.array([i * N + j for i, j in itertools.product(range(N), range(N))], dtype=np.intp)
test = Test(a, b, flat_indices)

这个版本的性能会非常接近run2()，因为：

• 只需要读取一次扁平化索引，减少了内存开销
• 使用直接指针访问数组，避免了Cython内存视图的索引计算开销
• 编译器可以对指针访问做更充分的优化

方案2：拆分索引为两个一维数组 #

如果必须保留二维索引的结构，可以将indices拆分为两个单独的一维数组（ixs和iys），这样内存访问更连续，同时Cython的优化会更高效：

修改后的Cython代码 #

%%cython -f -c=-O2 -I./
import numpy as np
cimport numpy as np
cimport cython

cdef class Test:
    cdef double[:, ::1] a, b
    cdef Py_ssize_t[::1] ixs, iys  # 拆分后的索引数组
    cdef Py_ssize_t ny

    def __cinit__(self, a, b, ixs, iys):
        self.a = a
        self.b = b
        self.ixs = ixs
        self.iys = iys
        self.ny = a.shape[1]

    @cython.boundscheck(False)
    @cython.nonecheck(False)
    @cython.wraparound(False)
    @cython.initializedcheck(False)
    cpdef void run_split(self):
        """使用拆分的一维索引数组"""
        cdef Py_ssize_t idx, ix, iy
        cdef int n = self.ixs.shape[0]
        cdef double* a_ptr = &self.a[0, 0]
        cdef double* b_ptr = &self.b[0, 0]

        for idx in range(n):
            ix = self.ixs[idx]
            iy = self.iys[idx]
            b_ptr[ix * self.ny + iy] = ix * iy * a_ptr[ix * self.ny + iy]

Python侧调用代码 #

import itertools
import numpy as np
N = 256
a = np.random.rand(N, N)
b = np.zeros_like(a)
indices = np.array([[i, j] for i, j in itertools.product(range(N), range(N))], dtype=np.intp)
# 拆分为两个一维数组
ixs = indices[:, 0]
iys = indices[:, 1]
test = Test(a, b, ixs, iys)

这个方案的性能也会比原始的run1()好很多，因为拆分后的一维数组内存访问更连续，Cython可以更好地优化缓存预取。

方案3：编译时优化索引访问 #

如果你能确保indices是连续的C顺序数组，可以在Cython中直接用指针访问indices的内存，减少数组索引的开销：

修改后的run1()优化版 #

@cython.boundscheck(False)
@cython.nonecheck(False)
@cython.wraparound(False)
@cython.initializedcheck(False)
cpdef void run1_optimized(self):
    """优化后的外部索引访问"""
    cdef Py_ssize_t idx, ix, iy
    cdef int n = self.indices.shape[0]
    cdef Py_ssize_t* indices_ptr = &self.indices[0, 0]  # 指向indices的首地址
    cdef double* a_ptr = &self.a[0, 0]
    cdef double* b_ptr = &self.b[0, 0]
    cdef Py_ssize_t ny = self.a.shape[1]

    for idx in range(n):
        # 因为indices是C顺序，所以indices[idx,0] = indices_ptr[2*idx], indices[idx,1] = indices_ptr[2*idx+1]
        ix = indices_ptr[2 * idx]
        iy = indices_ptr[2 * idx + 1]
        b_ptr[ix * ny + iy] = ix * iy * a_ptr[ix * ny + iy]

这个版本直接用指针访问indices的内存，避免了Cython内存视图的二维索引计算开销，性能会比原始run1()提升不少。

使用上述方案后，run_flat()或run_split()的性能应该会和run2()非常接近（误差在5%以内），具体取决于编译器的优化程度。

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