在Cython中使用并行化但性能无提升的问题排查

问题重现 #

在Windows环境的Jupyter Notebook中，按照Cython官方教程的「面向NumPy用户的Cython」章节实现多线程示例时，无论设置num_threads=4还是num_threads=1，程序运行速度均无变化，多线程未发挥作用。测试代码如下：

%%cython --force

import numpy as np
import timeit
cimport cython
from cython.parallel import prange

array_1 = np.random.uniform(0, 100, size=(3000, 2000)).astype(np.intc)
array_2 = np.random.uniform(0, 100, size=(3000, 2000)).astype(np.intc)
a = 4
b = 3
c = 9

ctypedef fused my_type:
    int
    double
    long long

cdef my_type clip(my_type a, my_type min_value, my_type max_value) nogil:
    return min(max(a, min_value), max_value)


@cython.boundscheck(False)
@cython.wraparound(False)
def compute_cy_mt(my_type[:, ::1] array_1, my_type[:, ::1] array_2, my_type a, my_type b, my_type c):

    cdef Py_ssize_t x_max = array_1.shape[0]
    cdef Py_ssize_t y_max = array_1.shape[1]

    assert tuple(array_1.shape) == tuple(array_2.shape)

    if my_type is int:
        dtype = np.intc
    elif my_type is double:
        dtype = np.double
    elif my_type is cython.longlong:
        dtype = np.longlong

    result = np.zeros((x_max, y_max), dtype=dtype)
    cdef my_type[:, ::1] result_view = result

    cdef my_type tmp
    cdef Py_ssize_t x, y

        for y in prange(y_max, num_threads=4,nogil=True):
            tmp = clip(array_1[x, y], 2, 10)
            tmp = tmp * a + array_2[x, y] * b
            result_view[x, y] = tmp + c

    return result

print(compute_cy_mt(array_1, array_2, a, b, c))

问题原因及解决方案 #

1. 修复循环逻辑错误 #

代码中仅对y维度进行循环，但未遍历x维度，导致实际只处理了未初始化的x对应的单行数据，计算量极小，多线程无法体现性能优势。需要修正为遍历所有元素，将prange应用到外层循环（比如x维度），内层循环处理y维度：

# 修正后的循环部分
for x in prange(x_max, num_threads=4, nogil=True):
    for y in range(y_max):
        tmp = clip(array_1[x, y], 2, 10)
        tmp = tmp * a + array_2[x, y] * b
        result_view[x, y] = tmp + c

2. Windows环境启用OpenMP编译 #

Cython的prange依赖OpenMP实现多线程，Windows下使用MSVC编译器时，需要在编译时显式添加OpenMP选项。在Jupyter的%%cython魔法命令中，需添加--compile-args=/openmp参数：

%%cython --force --compile-args=/openmp

完整修正后的代码 #

%%cython --force --compile-args=/openmp

import numpy as np
import timeit
cimport cython
from cython.parallel import prange

array_1 = np.random.uniform(0, 100, size=(3000, 2000)).astype(np.intc)
array_2 = np.random.uniform(0, 100, size=(3000, 2000)).astype(np.intc)
a = 4
b = 3
c = 9

ctypedef fused my_type:
    int
    double
    long long

cdef my_type clip(my_type a, my_type min_value, my_type max_value) nogil:
    return min(max(a, min_value), max_value)


@cython.boundscheck(False)
@cython.wraparound(False)
def compute_cy_mt(my_type[:, ::1] array_1, my_type[:, ::1] array_2, my_type a, my_type b, my_type c):

    cdef Py_ssize_t x_max = array_1.shape[0]
    cdef Py_ssize_t y_max = array_1.shape[1]

    assert tuple(array_1.shape) == tuple(array_2.shape)

    if my_type is int:
        dtype = np.intc
    elif my_type is double:
        dtype = np.double
    elif my_type is cython.longlong:
        dtype = np.longlong

    result = np.zeros((x_max, y_max), dtype=dtype)
    cdef my_type[:, ::1] result_view = result

    cdef my_type tmp
    cdef Py_ssize_t x, y

    # 修正循环逻辑，使用prange并行处理x维度
    for x in prange(x_max, num_threads=4, nogil=True):
        for y in range(y_max):
            tmp = clip(array_1[x, y], 2, 10)
            tmp = tmp * a + array_2[x, y] * b
            result_view[x, y] = tmp + c

    return result

# 测试性能对比
print("4线程耗时:", timeit.timeit(lambda: compute_cy_mt(array_1, array_2, a, b, c), number=5))
# 修改num_threads为1后再次运行测试，对比耗时差异

额外验证步骤 #

• 确认安装的Cython版本为0.29及以上（该版本对OpenMP支持更完善）
• 检查Windows环境的MSVC编译器版本（VS2015及以上版本默认支持OpenMP）

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