与 Numpy 相比优化 Cython 循环

Question

#cython: boundscheck=False, wraparound=False, nonecheck=False, cdivision=True, language_level=3
cpdef int query(double[::1] q, double[:,::1] data) nogil:
    cdef:
        int n = data.shape[0]
        int dim = data.shape[1]
        int best_i = -1
        double best_ip = -1
        double ip
    for i in range(n):
        ip = 0
        for j in range(dim):
            ip += q[j] * data[i, j]
        if ip > best_ip:
            best_i = i
            best_ip = ip
    return best_i

编译后，我把代码从Python:

import numpy as np
import ip
n, dim = 10**6, 10**2
X = np.random.randn(n, dim)
q = np.random.randn(dim)
%timeit ip.query(q, X)

这大约需要 100 毫秒。同时相当于 numpy code:

%timeit np.argmax(q @ X.T)

只需要大约 50 毫秒。

这很奇怪，因为 NumPy 代码似乎必须在获取 argmax 之前分配大数组 q @ X.T。因此，我想知道我是否缺少一些优化？

我已将 extra_compile_args=["-O3", '-march=native'], 添加到我的 setup.py 并且我还尝试将函数定义更改为

cpdef int query(np.ndarray[double] q, np.ndarray[double, ndim=2] data):

但它在性能上几乎没有差异。

Answer 1

操作 q @ X.T 将映射到引擎盖下来自 OpenBlas 或 MKL（取决于您的分布）的矩阵向量乘法 (dgemv) 的实现 - 这意味着您是反对最好的优化算法之一。

生成的向量有 1M 个元素，这导致大约 8MB 内存。 8MB 并不总是适合 L3 缓存，但即使 RAM 也有大约 15GB/s 的带宽，因此 writing/reading 8MB 最多需要 1-2ms - 与整体的大约 50ms 相比增益不大运行时间。

您的代码最明显的问题是它的计算方式与 q @X.T 不同。它计算

((q[0]*data[i,0]+q[1]*data[i,1])+q[2]*data[i,2])+...

由于 IEEE 754，不允许编译器对操作重新排序并以这种非最佳顺序执行它们：为了计算第二个和，操作必须等到执行第一个求和。这种方法没有充分利用现代架构的潜力。

一个好的 dgemv 实现会选择一个更好的操作顺序。可以在 .

中找到类似的问题，但有总和

为了平衡这一领域，可以使用 -ffast-math，它允许编译器重新排序操作，从而更好地利用管道。

以下是我机器上的基准测试结果：

%timeit query(q, X)            # 101 ms
%timeit query_ffastmath(q, X)  # 56.3 ms
%timeit np.argmax(q @ X.T)     # 50.2 ms

它仍然差了大约 10%，但如果编译器能够击败由专家特别为我的处理器创建的手工版本，我会感到非常惊讶。

与 Numpy 相比优化 Cython 循环

Optimizing Cython loop compared to Numpy

python

optimization

numpy

cython