从BLAS1到BLAS3：一个矩阵乘法的性能优化之旅（附Python/C++代码对比）

在深度学习训练和科学计算的实践中，矩阵乘法往往是性能瓶颈所在。当我在优化一个自定义神经网络层时，最初用朴素的Python循环实现，发现处理500x500矩阵需要近10秒——这促使我踏上了探索BLAS分级优化的旅程。本文将带你从最基础的BLAS1级实现出发，逐步引入向量化、分块等技术，最终过渡到BLAS3的GEMM优化，并通过实测数据揭示不同级别实现的性能差异。

1. 性能优化的起点：BLAS1级实现

我们先从最直观的三重循环实现开始。这种实现对应BLAS1的思想——每次只计算一个元素，完全不考虑内存访问模式或CPU特性：

python复制def naive_matrix_mult(A, B):
    m, n = A.shape
    p = B.shape[1]
    C = np.zeros((m, p))
    for i in range(m):
        for j in range(p):
            for k in range(n):
                C[i,j] += A[i,k] * B[k,j]
    return C

在Intel i7-11800H上测试512x512矩阵乘法，这个实现耗时约9.8秒。性能低下的主要原因有三：

1. 内存访问模式差：B矩阵被按列访问，破坏了空间局部性
2. 计算强度低：每次乘加操作都需要两次内存访问
3. 未利用向量化：标量运算无法发挥CPU的SIMD能力

注：所有测试均在相同环境下进行，使用Python 3.9和C++20，编译器优化选项为-O3

2. 第一次飞跃：BLAS2级优化策略

BLAS2级操作引入了矩阵-向量运算的思想。我们首先优化内存访问模式：

python复制def improved_matrix_mult(A, B):
    m, n = A.shape
    p = B.shape[1]
    C = np.zeros((m, p))
    for i in range(m):
        for k in range(n):
            r = A[i,k