Cython 循环遍历索引数组答案

【问题标题】：Cython loop over array of indexesCython 循环遍历索引数组
【发布时间】：2022-06-30 03:06:22
【问题描述】：

我想对矩阵的特定元素进行一系列操作。我需要在外部对象中定义这些元素的索引（下例中的self.indices）。

这是一个在 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] indices
    
    def __cinit__(self, a, b, indices):
        self.a = a
        self.b = b
        self.indices = indices
    
    @cython.boundscheck(False)
    @cython.nonecheck(False)
    @cython.wraparound(False)
    @cython.initializedcheck(False)
    cpdef void run1(self):
        """ Use of external structure of indices. """
        cdef Py_ssize_t idx, ix, iy
        cdef int n = self.indices.shape[0]
        
        
        for idx in range(n):
            ix = self.indices[idx, 0]
            iy = self.indices[idx, 1]
            self.b[ix, iy] = ix * iy * self.a[ix, iy]

    @cython.boundscheck(False)
    @cython.nonecheck(False)
    @cython.wraparound(False)
    @cython.initializedcheck(False)
    cpdef void run2(self):
        """ Direct formulation """
        cdef Py_ssize_t idx, ix, iy
        cdef int nx = self.a.shape[0]
        cdef int ny = self.a.shape[1]
        
        for ix in range(nx):
            for iy in range(ny):
                self.b[ix, iy] = ix * iy * self.a[ix, 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=int)
test = Test(a, b, indices)

结果：

%timeit test.run1()
75.6 µs ± 1.51 µs per loop (mean ± std. dev. of 7 runs, 10,000 loops each)

%timeit test.run2()
41.4 µs ± 1.77 µs per loop (mean ± std. dev. of 7 runs, 10,000 loops each)

为什么Test.run1() 方法的运行速度比Test.run2() 方法慢很多？

通过使用外部 list、array 或任何其他类型的索引结构，有哪些可能保持与 Test.run2() 类似的性能水平？

【问题讨论】：

标签： python cython

【解决方案1】：

因为run1 要复杂得多...

run1 不得不从两个单独的位中读取更多的内存，这几乎肯定会降低 CPU 缓存的效率。
编译器可以很简单地计算出它访问run2 中数组元素的确切顺序。相比之下，run1 可以以任何顺序访问它们。这可能会进行重大优化。

您目前的表现可能已经达到了最佳水平。

【讨论】：