Python C 扩展比 Numba JIT 更快吗？

Question

我正在测试 Numba JIT 与 Python C 扩展的性能。对于基于 for 循环的函数来计算 2d 数组中所有元素的总和，C 扩展似乎比 Numba 等效项快 3-4 倍。

更新：

根据宝贵的意见，我意识到我应该编译（调用）一次 Numba JIT 的错误。我提供了修复后的测试结果以及额外的案例。但问题仍然是何时以及如何考虑哪种方法。

这是结果（time_s，值）：

# 200 tests mean (including JIT compile inside the loop)
Pure Python: (0.09232537984848023, 29693825)
Numba: (0.003188209533691406, 29693825)
C Extension: (0.000905141830444336, 29693825.0)

# JIT once called before the test loop (to avoid compile time)
Normal: (0.0948486328125, 29685065)
Numba: (0.00031280517578125, 29685065)
C Extension: (0.0025129318237304688, 29685065.0)

# JIT no warm-up also no test loop (only calling once)
Normal: (0.10458517074584961, 29715115)
Numba: (0.314251184463501, 29715115)
C Extension: (0.0025091171264648438, 29715115.0)

• 我的实现正确吗？
• C 扩展更快有什么原因吗？
• 如果我想要最佳性能，我是否应该始终使用 C 扩展？（非向量化函数）

main.py

# 200 tests mean (including JIT compile inside the loop)
Pure Python: (0.09232537984848023, 29693825)
Numba: (0.003188209533691406, 29693825)
C Extension: (0.000905141830444336, 29693825.0)

# JIT once called before the test loop (to avoid compile time)
Normal: (0.0948486328125, 29685065)
Numba: (0.00031280517578125, 29685065)
C Extension: (0.0025129318237304688, 29685065.0)

# JIT no warm-up also no test loop (only calling once)
Normal: (0.10458517074584961, 29715115)
Numba: (0.314251184463501, 29715115)
C Extension: (0.0025091171264648438, 29715115.0)

ext.c

import numpy as np
import pandas as pd
import numba
import time
import loop_test # ext


def test(fn, *args):
    res = []
    val = None
    for _ in range(100):
        start = time.time()
        val = fn(*args)
        res.append(time.time() - start)
    return np.mean(res), val


sh = (30_000, 20)
col_names = [f"col_{i}" for i in range(sh[1])]
df = pd.DataFrame(np.random.randint(0, 100, size=sh), columns=col_names)
arr = df.to_numpy()


def sum_columns(arr):
    _sum = 0
    for i in range(arr.shape[0]):
        for j in range(arr.shape[1]):
            _sum += arr[i, j]
    return _sum


@numba.njit
def sum_columns_numba(arr):
    _sum = 0
    for i in range(arr.shape[0]):
        for j in range(arr.shape[1]):
            _sum += arr[i, j]
    return _sum


print("Pure Python:", test(sum_columns, arr))
print("Numba:", test(sum_columns_numba, arr))
print("C Extension:", test(loop_test.loop_fn, arr))

setup.py

#define PY_SSIZE_CLEAN
#include <Python.h>
#include <numpy/arrayobject.h>

static PyObject *loop_fn(PyObject *module, PyObject *args)
{
    PyObject *arr;
    if (!PyArg_ParseTuple(args, "O!", &PyArray_Type, &arr))
        return NULL;

    npy_intp *dims = PyArray_DIMS(arr);
    npy_intp rows = dims[0];
    npy_intp cols = dims[1];
    double sum = 0;
    PyArrayObject *arr_new = (PyArrayObject *)PyArray_FROM_OTF(arr, NPY_DOUBLE, NPY_ARRAY_IN_ARRAY);
    double *data = (double *)PyArray_DATA(arr_new);
    npy_intp i, j;
    for (i = 0; i < rows; i++)
        for (j = 0; j < cols; j++)
            sum += data[i * cols + j];
    Py_DECREF(arr_new);
    return Py_BuildValue("d", sum);
};

static PyMethodDef Methods[] = {
    {
        .ml_name = "loop_fn",
        .ml_meth = loop_fn,
        .ml_flags = METH_VARARGS,
        .ml_doc = "Returns the sum using for loop, but in C.",
    },
    {NULL, NULL, 0, NULL},
};

static struct PyModuleDef Module = {
    PyModuleDef_HEAD_INIT,
    "loop_test",
    "A benchmark module test",
    -1,
    Methods};

PyMODINIT_FUNC PyInit_loop_test(void)
{
    import_array();
    return PyModule_Create(&Module);
}

python3 setup.py install

Answer 1

我想完成约翰·博林格（John Bollinger）的精彩回答：

首先，C 扩展倾向于在 Linux 上使用 GCC 进行编译（可能是 Windows 上的 MSVC 和 MacOS 上的 Clang AFAIK），而Numba 在内部使用 LLVM编译工具链。如果你想比较两者，那么你应该使用基于 LLVM 工具链的 Clang。事实上，您还应该使用与 Numba 相同版本的 LLVM，以便进行公平比较。Clang、GCC 和 MSVC 优化代码的方式不同，因此生成的程序可能具有截然不同的性能。

此外，Numba 是一种 JIT，因此它不关心不同平台之间（指令集扩展）的兼容性。这意味着它可以使用 AVX-2 SIMD 指令集（如果您的机器上可用），而主流编译器出于兼容性考虑默认不会这样做。事实上，Numba 确实这么做了。您可以指定 Clang 和 GCC 来优化目标机器的代码，而不用关心带有编译标志的机器之间的兼容性-march=native。因此，生成的包肯定会更快，但也可能在旧机器上崩溃（或者可能显着变慢）。您还可以启用一些特定的指令集（带有类似的标志-mavx2）。

此外， Numba默认情况下使用积极的优化级别-O2，而 AFAIK C 扩展使用的标志默认情况下不会在 GCC 和 Clang 上自动矢量化代码（即不使用打包的 SIMD 指令）。-O3如果尚未完成，您当然应该手动指定使用该标志。在 MSVC 上，等效标志是/O2（据我所知还没有/O3）。

请注意，可以通过提供特定签名（可能是多个签名）来快速编译 Numba 函数（而不是默认情况下延迟编译）。这意味着您应该知道输入参数的类型，并且应用程序的启动时间可能会显着增加。Numba 函数也可以被缓存，这样就不会在同一平台上一遍又一遍地重新编译该函数。这可以通过标志来完成cache=True。但对于您的特定用例来说，它可能并不总是有效。

最后但并非最不重要的一点是，这两个代码并不等效。这当然是最重要的一点。Numba 代码处理 类型int32并将arr值累加为 64 位整数_sum，而 C 扩展则将值累加为double精度浮点类型。浮点类型不是关联的（除非您使用 flag 告诉编译器假设它们是关联的，-ffast-math该标志默认情况下不启用，因为它不安全），因此由于高延迟，累加浮点数比整数昂贵得多大多数平台上的 FMA 单位。此外，我实际上想知道是否PyArray_FROM_OTF执行了正确的转换，但如果确实如此，那么我预计转换会非常昂贵。为了公平比较，您应该在两个代码中使用相同的类型（两个代码中可能是 64 位整数）。

欲了解更多信息，请阅读相关帖子：

• 在 Cython 中循环二维数组的最快方法
• 对于这个简单的循环，为什么 cython 比 numba 慢得多？