Python无GPU场景下加速稀疏B与稠密A的B.T@A@B矩阵运算的方案

1 利用矩阵数学性质降低计算量 #

• 优先调整运算顺序：利用矩阵乘法结合律，先计算A @ B再和B.T相乘，同时因为A是对称矩阵，主流BLAS库会自动对对称矩阵乘稠密/稀疏矩阵的运算做减半优化，比直接计算(B.T @ A) @ B效率提升至少1倍
• 利用结果对称性：最终输出是2000×2000的对称矩阵，只需计算上三角/下三角部分的元素，剩余部分直接赋值复用，可降低近一半的最终内积计算量
• 利用B的稀疏性：将运算拆解为B列向量之间的内积计算b_i^T A b_j，跳过B中零元素的冗余计算，稀疏度越高收益越明显

2 运算库与存储格式优化 #

• 采用MKL/OpenBLAS优化的NumPy、SciPy库：Intel MKL对对称矩阵运算、稀疏矩阵乘稠密矩阵的运算做了深度指令集优化，Intel CPU下相比原生NumPy性能可提升2~4倍
• 选择合适的稀疏矩阵存储格式：将B转换为scipy.sparse.csc_matrix格式（按列存储），取列向量的效率远高于COO等格式，适合频繁访问B列的计算场景
• 数据类型降精度优化：如果业务允许精度损失，将A、B统一转换为float32类型，内存占用直接减半，同时CPU SIMD指令可并行处理更多单精度运算，整体速度可提升50%以上

3 手动实现与并行优化 #

• 用Numba JIT编译自定义计算逻辑：避免Python原生循环的开销，搭配prange实现多线程并行计算不同的(i,j)元素对，示例代码框架如下：

import numpy as np
from numba import jit, prange
from scipy.sparse import csc_matrix

@jit(nopython=True, parallel=True)
def compute_product(A, B_data, B_indices, B_indptr, k, n):
    res = np.zeros((k, k), dtype=A.dtype)
    for i in prange(k):
        # 取出B的第i列非零元素
        start_i = B_indptr[i]
        end_i = B_indptr[i+1]
        idx_i = B_indices[start_i:end_i]
        val_i = B_data[start_i:end_i]
        # 计算A @ b_i的有效部分
        Ab_i = np.zeros(n, dtype=A.dtype)
        for pos in range(len(idx_i)):
            row = idx_i[pos]
            Ab_i += A[:, row] * val_i[pos]
        # 计算上三角元素
        for j in range(i, k):
            start_j = B_indptr[j]
            end_j = B_indptr[j+1]
            idx_j = B_indices[start_j:end_j]
            val_j = B_data[start_j:end_j]
            tmp = 0.0
            for pos in range(len(idx_j)):
                row = idx_j[pos]
                tmp += Ab_i[row] * val_j[pos]
            res[i, j] = tmp
            res[j, i] = tmp
    return res

# 调用示例
B = csc_matrix(B)
k = B.shape[1]
n = B.shape[0]
result = compute_product(A, B.data, B.indices, B.indptr, k, n)

• 合理配置BLAS线程数：提前设置环境变量OMP_NUM_THREADS为CPU物理核心数，避免超线程带来的性能损耗，例如16物理核可设置os.environ['OMP_NUM_THREADS'] = '16'

4 数值近似方案（精度允许时使用） #

• 对称矩阵低秩近似：对A做特征分解取前m个最大特征值对应的特征向量（m可根据精度要求取1000~5000），将A近似为低秩形式后再计算乘积，计算量可降低一个数量级
• Nyström近似：如果A是对称正定矩阵，采用Nyström方法采样A的部分行列构建近似矩阵，最高可将整体运算速度提升10倍以上

5 硬件适配优化 #

• 确保内存容量足够：40000×40000的float64稠密矩阵A占12.8GB内存，中间结果A@B占640MB，需保证机器可用内存大于16GB，避免触发swap交换分区拖慢速度
• 内存不足时采用分块计算：如果无法加载完整A矩阵，可将A切分为多个小块并行运算，避免内存溢出
• 优先使用大缓存CPU：该运算为访存密集型，CPU缓存越大，访问A矩阵元素的命中率越高，运算速度越快

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