AVX矩阵乘法分块大小设置问题：块大小16时结果异常

问题描述 #

我编写了基于AVX的分块矩阵乘法代码，当块大小设为4时运行正确，但改为16时，输出矩阵仅1/4的结果正确。已知x86有16个YMM寄存器，疑惑为何块大小不能设为16。

原始代码 #

#include <x86intrin.h>
#include <immintrin.h>
#include <stdio.h>
#include <stdlib.h>
#include <time.h>

void fill(double *mat, int n) {
  register int i;
  register int j;
  for (i = 0; i < n; ++i) {
    for (j = 0; j < n; ++j) {
      mat[i * n + j] = (double) (i * n + j);
    }
  }
}

void blocked_mmul_avx(double *A, double *B, double *C, int n) {
  register int row, block, chunk, sub_chunk;
  int block_size = 4; // 修改为16时结果仅1/4正确

  for (row = 0; row < n; row++)
    for (block = 0; block < n; block += block_size) {
      register __m256d c0 = _mm256_setzero_pd();
      for (chunk = 0; chunk < n; chunk += block_size) {
        for (sub_chunk = 0; sub_chunk < block_size; sub_chunk++) {
          register __m256d a = _mm256_broadcast_sd(
              &A[row * n + chunk + sub_chunk]);
          register __m256d b = _mm256_loadu_pd(
              &B[chunk * n + sub_chunk * n + block]);
          c0 += a * b;
        }
      }
      _mm256_storeu_pd(&C[row * n + block], c0);
    }
}

int main() {
  int n = 4096; //n = 512, 1024, 2048, 4096

  double *A;
  double *B;
  double *C;

  A = (double*) malloc(n * n * sizeof(double));
  B = (double*) malloc(n * n * sizeof(double));
  C = (double*) malloc(n * n * sizeof(double));

  srand(0);
  fill(A, n);  // 用于为A、B加载数据的函数
  fill(B, n);

  printf("n = %d\n\n", n);
  clock_t t0 = clock();
  blocked_mmul_avx(A, B, C, n);
  clock_t t1 = clock();
  printf("Execution Time = %0.3f s \n\n", (float) (t1 - t0) / CLOCKS_PER_SEC);

  free(A);
  free(B);
  free(C);
  return 0;
}

原因分析 #

• 单寄存器容量限制：__m256d是256位向量寄存器，每个最多容纳4个double类型元素（8字节×4=32字节=256位）。你的代码仅使用c0这一个寄存器累加结果，当block_size=16时，你试图一次性处理16列的计算，但c0只能存储4列的结果，剩余12列的C矩阵位置从未被正确赋值，因此仅前4列（即1/4）结果正确。
• 寄存器数量与块大小的误解：16个YMM寄存器是处理器的硬件资源，但你的代码并未利用多寄存器并行处理多列结果。要支持block_size=16，需要对应使用4个__m256d寄存器（16÷4=4），每个寄存器负责4列的乘法累加计算。

修正示例（支持block_size=16） #

修改blocked_mmul_avx函数，使用4个累加寄存器分别处理16列中的每4列：

void blocked_mmul_avx(double *A, double *B, double *C, int n) {
  register int row, block, chunk, sub_chunk;
  int block_size = 16; // 现在支持16的块大小
  int vec_width = 4; // __m256d一次处理4个double

  for (row = 0; row < n; row++)
    for (block = 0; block < n; block += block_size) {
      // 初始化4个累加寄存器，对应16列中的4组
      __m256d c0 = _mm256_setzero_pd();
      __m256d c1 = _mm256_setzero_pd();
      __m256d c2 = _mm256_setzero_pd();
      __m256d c3 = _mm256_setzero_pd();

      for (chunk = 0; chunk < n; chunk += block_size) {
        for (sub_chunk = 0; sub_chunk < block_size; sub_chunk++) {
          __m256d a = _mm256_broadcast_sd(&A[row * n + chunk + sub_chunk]);
          // 加载B中对应的4组4列数据
          __m256d b0 = _mm256_loadu_pd(&B[(chunk + sub_chunk)*n + block]);
          __m256d b1 = _mm256_loadu_pd(&B[(chunk + sub_chunk)*n + block + vec_width]);
          __m256d b2 = _mm256_loadu_pd(&B[(chunk + sub_chunk)*n + block + vec_width*2]);
          __m256d b3 = _mm256_loadu_pd(&B[(chunk + sub_chunk)*n + block + vec_width*3]);
          // 分别累加
          c0 += a * b0;
          c1 += a * b1;
          c2 += a * b2;
          c3 += a * b3;
        }
      }
      // 将4个寄存器的结果分别存储到C的对应位置
      _mm256_storeu_pd(&C[row*n + block], c0);
      _mm256_storeu_pd(&C[row*n + block + vec_width], c1);
      _mm256_storeu_pd(&C[row*n + block + vec_width*2], c2);
      _mm256_storeu_pd(&C[row*n + block + vec_width*3], c3);
    }
}

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