Python C扩展是否比Numba JIT更快？附性能测试与疑问

我正在测试Numba JIT与Python C扩展的性能，针对计算二维数组元素和的循环函数，最初发现C扩展比Numba等效实现快约3-4倍。

更新说明 #

根据反馈，我意识到之前未提前编译（调用）Numba JIT的错误，现提供修复后的测试结果及额外场景，但仍有疑问待解答。

测试结果（耗时_s, 计算值） #

# 200次测试均值（循环内包含JIT编译）
Pure Python: (0.09232537984848023, 29693825)
Numba: (0.003188209533691406, 29693825)
C Extension: (0.000905141830444336, 29693825.0)

# 测试循环前提前调用JIT（避免编译耗时）
Normal: (0.0948486328125, 29685065)
Numba: (0.00031280517578125, 29685065)
C Extension: (0.0025129318237304688, 29685065.0)

# 无预热且仅调用一次（无测试循环）
Normal: (0.10458517074584961, 29715115)
Numba: (0.314251184463501, 29715115)
C Extension: (0.0025091171264648438, 29715115.0)

疑问 #

1. 我的实现是否正确？
2. C扩展更快的原因是什么？
3. 若追求最优性能（非向量化函数），是否应始终使用C扩展？

附测试代码 #

main.py #

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))

ext.c #

#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);
}

setup.py #

from distutils.core import setup, Extension
import numpy as np

module = Extension(
    "loop_test",
    sources=["ext.c"],
    include_dirs=[
        np.get_include(),
    ],
)

setup(
    name="loop_test",
    version="1.0",
    description="This is a test package",
    ext_modules=[module],
)

执行命令 #

python3 setup.py install

内容的提问来源于stack exchange，提问作者Momo