如何优化Cython矩阵轴循环代码以充分利用AVX2 SIMD指令？

针对你的Cython代码未充分利用SIMD指令的问题，以下是具体修改方案，从消除优化阻碍到强制SIMD利用，逐步提升性能：

1. 关闭Cython安全检查，消除优化障碍 #

Cython默认的边界检查、数组越界包装会干扰编译器的向量化判断，先关闭这些选项，并显式指定内存布局：

# cython: boundscheck=False, wraparound=False, cdivision=True, language_level=3
import numpy as np
from libc.math cimport fabs

cpdef inline double[:,:] go(double[:] stream, double[:] query):
    # 显式指定C行优先顺序，保证内存连续访问
    matrix = np.empty((len(stream), len(query)), order='C')
    cdef double [:, :] matrix_c = matrix
   
    cdef int i, j
    cdef int stream_len = len(stream)
    cdef int query_len = len(query)
    for i in range(stream_len):
        # 提前加载stream元素，减少重复内存访问
        cdef double s_val = stream[i]
        for j in range(query_len):
            matrix_c[i, j] = fabs(s_val - query[j])

    return matrix

编译时添加-ffast-math参数，允许编译器对浮点运算进行激进优化（包括向量化fabs）：

cythonize -i -O3 -march=native -ffast-math go.pyx

2. 替换标量函数为SIMD原生指令 #

直接调用AVX2内置函数，强制生成SIMD指令。AVX2的256位寄存器可同时处理4个double类型元素：

# cython: boundscheck=False, wraparound=False, cdivision=True, language_level=3
import numpy as np
cimport numpy as np

# 声明AVX2内置函数
cdef extern from "immintrin.h":
    np.float64_t[:] _mm256_loadu_pd(double* ptr) nogil
    np.float64_t[:] _mm256_set1_pd(double val) nogil
    np.float64_t[:] _mm256_sub_pd(np.float64_t[:] a, np.float64_t[:] b) nogil
    np.float64_t[:] _mm256_abs_pd(np.float64_t[:] a) nogil
    void _mm256_storeu_pd(double* ptr, np.float64_t[:] val) nogil

cpdef inline double[:,:] go(double[:] stream, double[:] query):
    cdef int stream_len = len(stream)
    cdef int query_len = len(query)
    # 补全query长度为4的倍数，适配AVX2向量宽度
    cdef int padded_len = ((query_len + 3) // 4) * 4
    cdef double[:] padded_query = np.zeros(padded_len, dtype=np.float64)
    padded_query[:query_len] = query
    
    matrix = np.empty((stream_len, query_len), order='C')
    cdef double [:, :] matrix_c = matrix
   
    cdef int i, j
    for i in range(stream_len):
        # 将单个stream元素广播为AVX2向量
        cdef np.float64_t[:] s_vec = _mm256_set1_pd(stream[i])
        # 按向量宽度批量处理query元素
        for j in range(0, padded_len, 4):
            cdef np.float64_t[:] q_vec = _mm256_loadu_pd(&padded_query[j])
            cdef np.float64_t[:] diff_vec = _mm256_sub_pd(s_vec, q_vec)
            cdef np.float64_t[:] abs_vec = _mm256_abs_pd(diff_vec)
            _mm256_storeu_pd(&matrix_c[i, j], abs_vec)
        # 处理剩余不足4个的元素
        for j in range(padded_len - (padded_len - query_len), query_len):
            matrix_c[i, j] = fabs(stream[i] - query[j])

    return matrix

编译时保持-march=native参数，确保编译器识别AVX2指令集。

3. 改用NumPy广播替代手动循环 #

NumPy底层已做高度优化，会自动利用SIMD指令处理广播运算，代码更简洁且性能优异：

# cython: language_level=3
import numpy as np

cpdef inline double[:,:] go(double[:] stream, double[:] query):
    # 将stream转为列向量，query转为行向量，广播后计算差值绝对值
    return np.abs(np.asarray(stream)[:, np.newaxis] - np.asarray(query)[np.newaxis, :])

4. 用OpenMP SIMD显式引导向量化 #

启用OpenMP并添加SIMD编译指令，强制编译器对循环进行向量化：

# cython: boundscheck=False, wraparound=False, cdivision=True, language_level=3, openmp=True
import numpy as np
from libc.math cimport fabs

cpdef inline double[:,:] go(double[:] stream, double[:] query):
    matrix = np.empty((len(stream), len(query)), order='C')
    cdef double [:, :] matrix_c = matrix
   
    cdef int i, j
    cdef int stream_len = len(stream)
    cdef int query_len = len(query)
    for i in range(stream_len):
        cdef double s_val = stream[i]
        # 显式指定内层循环进行SIMD向量化
        #pragma omp simd
        for j in range(query_len):
            matrix_c[i, j] = fabs(s_val - query[j])

    return matrix

编译时需添加OpenMP相关参数：

cythonize -i -O3 -march=native -ffast-math -fopenmp go.pyx

5. 调整循环结构提升缓存命中率 #

将stream分块处理，适配CPU缓存大小，减少缓存失效：

# cython: boundscheck=False, wraparound=False, cdivision=True, language_level=3
import numpy as np
from libc.math cimport fabs

cpdef inline double[:,:] go(double[:] stream, double[:] query):
    cdef int stream_len = len(stream)
    cdef int query_len = len(query)
    # 块大小适配L1缓存，可根据CPU调整
    cdef int block_size = 256
    matrix = np.empty((stream_len, query_len), order='C')
    cdef double [:, :] matrix_c = matrix
   
    cdef int i_block, i, j
    for i_block in range(0, stream_len, block_size):
        cdef int block_end = min(i_block + block_size, stream_len)
        for i in range(i_block, block_end):
            cdef double s_val = stream[i]
            for j in range(query_len):
                matrix_c[i, j] = fabs(s_val - query[j])

    return matrix

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