python - 与 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代码貌似要分配大数组 q @ X.T在使用 argmax 之前。因此我想知道我是否缺少一些优化？
我已添加 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 的带宽，因此写入/读取 8MB 最多需要 1-2 毫秒 - 与大约 50 毫秒的总体运行时间相比，没有多少增益。
您的代码最明显的问题是它的计算方式与 q @X.T 不同。 .它计算

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

由于 IEEE 754，编译器不允许重新排序操作并以这种非最佳顺序执行它们:为了计算第二个总和，操作必须等到执行第一个总和。这种方法没有充分利用现代架构的潜力。
不错 dgemv实现将选择更好的操作顺序。可以在此 SO-post 中找到类似的问题，但带有总和。 .
可以使用 -ffast-math 来平整 field ，它允许编译器重新编码操作，从而更好地利用管道。
以下是我的机器上的基准测试结果:

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

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