Cython实现速度远低于Numba接近纯Python的原因排查

问题背景 #

• 日常仿真程序通常使用Numba实现加速，但由于Numba不支持代码中需要调用的某款SciPy函数，必须将代码迁移至Cython实现。
• 完成Cython版本改写后，程序运行速度远低于同逻辑Numba版本，性能表现和原生Python几乎持平。
• 基准测试调用参数为simulation(N=1000000, N_save=10000)，测试结果如下：
  • Cython版本：183ms
  • Numba版本：31.4ms
  • 原生Python版本：217ms

相关代码 #

初始Cython实现 #

import numpy as np
cimport numpy as cnp
cimport cython

cnp.import_array()

@cython.boundscheck(False)
@cython.wraparound(False)
def simulation(int N=10000, int N_save=10000):

    cdef cnp.ndarray[double] x = np.empty(N_save, dtype=np.double)
    cdef cnp.ndarray[double] r = np.random.standard_normal(size=N_save)
    
    cdef fs = int(N / N_save)
    cdef xold = 0
    x[0] = 0

    for i in range(1, N):

        if i%N_save == 0:
            r = np.random.standard_normal(size=N_save)

        xnew = xold + r[i%N_save]

        if (i % fs) == 0:
            x[int(i / fs)] = xnew

        xold = xnew

    return x

Cython编译配置 #

from setuptools import setup
from Cython.Build import cythonize
import numpy as np

setup(
    ext_modules=cythonize(
        "test.pyx",
        compiler_directives={'language_level' : "3"},
    ), 
    include_dirs=[np.get_include()]
)

同逻辑Numba实现 #

import numba as nb
import numpy as np

@nb.jit(nopython=True, fastmath=True)
def simulation_nb(N=10000, N_save=10000):

    x = np.zeros(N_save)
    r = np.random.standard_normal(size=N_save)
    
    fs = int(N / N_save)
    xold = 0
    x[0] = 0


    for i in range(1, N):

        if i%N_save == 0:
            r = np.random.standard_normal(size=N_save)

        xnew = xold + r[i%N_save]

        if (i % fs) == 0:
            x[int(i / fs)] = xnew

        xold = xnew

    return x

性能差的核心原因 #

你的Cython代码只做了最表层的类型声明，核心循环几乎没有被编译成原生C代码，绝大多数操作仍然在走Python对象调度逻辑，所以速度和纯Python接近：

1. 变量类型缺失：fs、xold、xnew、循环内的索引值都没有明确声明为C原生类型，Cython默认将其处理为Python对象，所有算术运算、赋值操作都要走Python对象的内存分配、类型检查逻辑，开销是原生C运算的数十倍。
2. 数组访问效率低：声明numpy数组时没有指定维度，也没有使用Cython推荐的类型化内存视图（memoryview），每次数组索引都要调用numpy的Python层API，无法直接访问数组的底层C内存。
3. 运算优化开关未开全：没有开启cdivision选项，整数除法、取模运算都会保留Python的溢出、除零检查逻辑，拖慢运算速度。
4. 随机数调用开销高：循环内调用np.random.standard_normal是Python层的函数调用，每次调用都有Python栈调度开销；Numba在nopython模式下会把随机数生成逻辑完全编译为原生机器码，没有这部分额外开销。
5. 编译器优化未开启：编译配置没有加O3级优化、本机架构指令集适配参数，编译器默认的低优化级别无法生成高效机器码。

优化方案 #

按照Cython静态类型规范补全所有类型声明、开启对应优化选项即可达到和Numba相当的性能，优化后的参考代码：

import numpy as np
cimport numpy as cnp
cimport cython

cnp.import_array()

@cython.boundscheck(False)
@cython.wraparound(False)
@cython.cdivision(True)  # 关闭C运算的Python异常检查
def simulation(int N=10000, int N_save=10000):
    # 用一维连续内存视图声明数组，支持直接访问底层内存
    cdef double[::1] x = np.empty(N_save, dtype=np.double)
    cdef double[::1] r
    cdef int fs = N // N_save
    cdef double xold = 0.0
    cdef double xnew
    cdef int i, r_idx, x_idx

    # 初始化随机数数组
    r = np.random.standard_normal(size=N_save)
    x[0] = 0.0

    for i in range(1, N):
        r_idx = i % N_save
        if r_idx == 0:
            r = np.random.standard_normal(size=N_save)

        xnew = xold + r[r_idx]

        if (i % fs) == 0:
            x_idx = i // fs
            x[x_idx] = xnew

        xold = xnew

    # 内存视图转numpy数组返回
    return np.asarray(x)

同时修改编译配置，开启编译器最高优化：

from setuptools import setup
from Cython.Build import cythonize
import numpy as np

setup(
    ext_modules=cythonize(
        "test.pyx",
        compiler_directives={
            'language_level' : "3",
            'boundscheck': False,
            'wraparound': False,
            'cdivision': True
        },
        extra_compile_args=["-O3", "-march=native"]  # 开O3优化，适配本机CPU指令集
    ), 
    include_dirs=[np.get_include()]
)

如果需要进一步压缩随机数生成的开销，可以替换为C标准库的随机数生成器或者numpy的C级随机数API，性能还能再提升15%-20%。

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