从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]
            for j in range(p):
                C[i,j] += r * B[k,j]
    return C

这个版本通过循环交换，使B矩阵的访问变为连续模式。实测性能提升到3.2秒，但仍存在以下问题：

C++实现可以进一步利用SIMD指令：

cpp复制void matrix_mult_blas2(float* A, float* B, float* C, int m, int n, int p) {
    #pragma omp parallel for
    for (int i = 0; i < m; ++i) {
        for (int k = 0; k < n; ++k) {
            __m256 a = _mm256_set1_ps(A[i*n + k]);
            for (int j = 0; j < p; j += 8) {
                __m256 b = _mm256_load_ps(&B[k*p + j]);
                __m256 c = _mm256_load_ps(&C[i*p + j]);
                c = _mm256_fmadd_ps(a, b, c);
                _mm256_store_ps(&C[i*p + j], c);
            }
        }
    }
}

这个版本将性能提升至0.8秒，主要得益于：

• AVX2指令集的256位向量化（8个float并行计算）
• OpenMP多线程并行
• 更好的寄存器重用

3. 突破性能瓶颈：BLAS3级分块技术

BLAS3的核心思想是通过分块（Tiling）提高计算强度。我们先看Python的概念实现：

python复制def blocked_matrix_mult(A, B, block_size=64):
    m, n = A.shape
    p = B.shape[1]
    C = np.zeros((m, p))
    for ii in range(0, m, block_size):
        for kk in range(0, n, block_size):
            for jj in range(0, p, block_size):
                for i in range(ii, min(ii+block_size, m)):
                    for k in range(kk, min(kk+block_size, n)):
                        a = A[i,k]
                        for j in range(jj, min(jj+block_size, p)):
                            C[i,j] += a * B[k,j]
    return C

分块大小的选择对性能影响巨大。以下是不同分块大小的性能对比：

C++的完整BLAS3实现需要考虑更多优化技巧：

cpp复制void gemm_blocked(float* A, float* B, float* C, int m, int n, int p) {
    const int BLOCK = 64;
    #pragma omp parallel for collapse(2)
    for (int ii = 0; ii < m; ii += BLOCK) {
        for (int jj = 0; jj < p; jj += BLOCK) {
            float buffer[BLOCK][BLOCK] = {0};
            for (int kk = 0; kk < n; kk += BLOCK) {
                for (int i = ii; i < min(ii+BLOCK, m); ++i) {
                    for (int k = kk; k < min(kk+BLOCK, n); ++k) {
                        __m256 a = _mm256_set1_ps(A[i*n + k]);
                        for (int j = jj; j < min(jj+BLOCK, p); j += 8) {
                            __m256 b = _mm256_load_ps(&B[k*p + j]);
                            __m256 c = _mm256_load_ps(&buffer[i-ii][j-jj]);
                            c = _mm256_fmadd_ps(a, b, c);
                            _mm256_store_ps(&buffer[i-ii][j-jj], c);
                        }
                    }
                }
            }
            // Write back to C
            for (int i = ii; i < min(ii+BLOCK, m); ++i) {
                for (int j = jj; j < min(jj+BLOCK, p); j += 8) {
                    __m256 c = _mm256_load_ps(&C[i*p + j]);
                    __m256 buf = _mm256_load_ps(&buffer[i-ii][j-jj]);
                    c = _mm256_add_ps(c, buf);
                    _mm256_store_ps(&C[i*p + j], c);
                }
            }
        }
    }
}

这个实现包含多个关键优化：

1. 三级缓存阻塞：针对L1/L2/L3缓存分别优化分块大小
2. 寄存器分块：使用buffer数组减少内存写入
3. 向量化+并行：结合AVX2和OpenMP
4. 循环展开：编译器辅助的循环展开优化

实测性能达到0.15秒，接近OpenBLAS的GEMM性能（0.12秒）。

4. 现代BLAS库的魔法：GEMM深度解析

现代BLAS库如OpenBLAS的GEMM实现融合了更多高级技巧：

分层优化策略：

1. 微内核（Micro Kernel）：手写汇编优化最内层循环

   • 使用FMA指令
   • 精心安排寄存器使用
   • 避免流水线停顿

2. 宏内核（Macro Kernel）：

   c复制for (int k = 0; k < K; k += KC) {
       pack_A(A + k*M, lda, bufferA);
       for (int j = 0; j < N; j += NC) {
           pack_B(B + j*K, ldb, bufferB);
           micro_kernel(bufferA, bufferB, C + j*M, ldc);
       }
   }
   

3. 并行策略：

   • 基于OpenMP的任务划分
   • NUMA感知的内存分配
   • 动态负载均衡

性能对比表：

注：峰值性能基于i7-11800H的94 GFLOPS理论值，超频后可能超过理论峰值

在实际项目中，我发现在以下场景需要特别注意：

• 小矩阵（<64x64）直接调用BLAS可能因启动开销反而更慢
• 非对齐内存访问会导致性能下降30%以上
• 混合精度计算时要注意保持计算一致性