为什么 numba 这么快？

Question

我想编写一个函数，它将采用lefts形状索引(N_ROWS,)我想编写一个函数，它将创建一个矩阵out = (N_ROWS, N_COLS)矩阵，使得out[i, j] = 1当且仅当j >= lefts[i]。在循环中执行此操作的一个简单示例如下：

class Looped(Strategy):
    def copy(self, lefts):
        out = np.zeros([N_ROWS, N_COLS])
        for k, l in enumerate(lefts): 
            out[k, l:] = 1
        return out

现在我希望它尽可能快，所以我对此函数有不同的实现：

1. 简单的Python循环
2. numba 与@njit
3. 我调用的纯 C++ 实现ctypes

以下是 100 次运行的平均结果：

Looped took 0.0011599776260009093
Numba took 8.886413300206186e-05
CPP took 0.00013200821400096175

因此 numba 大约是第二快的实现（即 C++ 实现）的 1.5 倍。我的问题是为什么？

• 我在类似的问题中听说过 cython 速度较慢，因为它没有使用设置的所有优化标志进行编译，但是编译 cpp 实现-O3是否足以让我拥有编译器将给我的所有可能的优化？
• 我不完全理解如何将 numpy 数组交给 c++，我是否无意中在此处复制了数据？

---

Looped took 0.0011599776260009093
Numba took 8.886413300206186e-05
CPP took 0.00013200821400096175

# numba implementation

@njit
def numba_copy(lefts):
    out = np.zeros((N_ROWS, N_COLS), dtype=np.float32)
    for k, l in enumerate(lefts): 
        out[k, l:] = 1.
    return out

    
class Numba(Strategy):
    def __init__(self) -> None:
        # avoid compilation time when timing 
        numba_copy(np.array([1]))

    def copy(self, lefts):
        return numba_copy(lefts)

// array copy cpp

extern "C" void copy(const long *lefts,  float *outdatav, int n_rows, int n_cols) 
{   
    for (int i = 0; i < n_rows; i++) {
        for (int j = lefts[i]; j < n_cols; j++){
            outdatav[i*n_cols + j] = 1.;
        }
    }
}

// compiled to a .so using g++ -O3 -shared -o array_copy.so array_copy.cpp

---

带有时间等的完整代码：

# using cpp implementation

class CPP(Strategy):

    def __init__(self) -> None:
        lib = ctypes.cdll.LoadLibrary("./array_copy.so")
        fun = lib.copy
        fun.restype = None
        fun.argtypes = [
            ndpointer(ctypes.c_long, flags="C_CONTIGUOUS"),
            ndpointer(ctypes.c_float, flags="C_CONTIGUOUS"),
            ctypes.c_long,
            ctypes.c_long,
            ]
        self.fun = fun

    def copy(self, lefts):
        outdata = np.zeros((N_ROWS, N_COLS), dtype=np.float32, )
        self.fun(lefts, outdata, N_ROWS, N_COLS)
        return outdata

Answer 1

Numba 目前使用 LLVM-Lite 将代码高效地编译为二进制文件（在 Python 代码已转换为 LLVM 中间表示之后）。-O3该代码像使用带有标志和的 Clang 编译的 C++ 代码一样进行了优化-march=native。最后一个参数非常重要，因为它使 LLVM 能够在相对较新的 x86-64 处理器上使用更广泛的 SIMD 指令：AVX 和 AVX2（对于最新的 Intel 处理器可能是 AVX512）。否则，默认情况下，Clang 和 GCC 仅使用 SSE/SSE2 指令（因为向后兼容）。

另一个区别来自于 GCC 和 Numba 的 LLVM 代码之间的比较。Clang/LLVM 倾向于积极展开循环，而 GCC 通常不会。这对生成的程序有显着的性能影响。事实上，您可以看到从 Clang 生成的汇编代码：

使用 Clang（每个循环 128 个项目）：

.LBB0_7:
        vmovups ymmword ptr [r9 + 4*r8 - 480], ymm0
        vmovups ymmword ptr [r9 + 4*r8 - 448], ymm0
        vmovups ymmword ptr [r9 + 4*r8 - 416], ymm0
        vmovups ymmword ptr [r9 + 4*r8 - 384], ymm0
        vmovups ymmword ptr [r9 + 4*r8 - 352], ymm0
        vmovups ymmword ptr [r9 + 4*r8 - 320], ymm0
        vmovups ymmword ptr [r9 + 4*r8 - 288], ymm0
        vmovups ymmword ptr [r9 + 4*r8 - 256], ymm0
        vmovups ymmword ptr [r9 + 4*r8 - 224], ymm0
        vmovups ymmword ptr [r9 + 4*r8 - 192], ymm0
        vmovups ymmword ptr [r9 + 4*r8 - 160], ymm0
        vmovups ymmword ptr [r9 + 4*r8 - 128], ymm0
        vmovups ymmword ptr [r9 + 4*r8 - 96], ymm0
        vmovups ymmword ptr [r9 + 4*r8 - 64], ymm0
        vmovups ymmword ptr [r9 + 4*r8 - 32], ymm0
        vmovups ymmword ptr [r9 + 4*r8], ymm0
        sub     r8, -128
        add     rbp, 4
        jne     .LBB0_7

使用 GCC（每个循环 8 个项目）：

.L5:
        mov     rdx, rax
        vmovups YMMWORD PTR [rax], ymm0
        add     rax, 32
        cmp     rdx, rcx
        jne     .L5

因此，为了公平起见，您需要将Numba 代码与使用 Clang 和上述优化标志编译的 C++ 代码进行比较。

请注意，根据您的需求和最后一级处理器缓存的大小，您可以使用非临时存储（NT 存储）编写更快的特定于平台的 C++ 代码。NT 存储告诉处理器不要将数组存储在其缓存中。使用 NT 存储写入数据比在 RAM 中写入大型数组更快，但是如果数组可以放入缓存中，则在复制后读取存储的数组时速度可能会变慢（因为必须从 RAM 重新加载数组）。在您的情况（4 MiB 阵列）中，尚不清楚这是否会更快。