稠密矩阵乘法单核优化

最新推荐文章于 2025-09-17 13:10:22 发布

原创

最新推荐文章于 2025-09-17 13:10:22 发布 · 2k 阅读

9 ·

CC 4.0 BY-SA版权

文章标签：

#矩阵 #线性代数 #hpc #c语言

稠密矩阵乘法单核优化：Arm

编译选项
循环变换
消除指针别名
循环展开
内存对齐和简单Blocking
两级Blocking+转置重组

本次作业为优化单核上的稠密矩阵乘法。
以下各个优化方法的效果均是在它之前叠加的。

机器：鲲鹏920芯片，只用64个核中的一个
提供的矩阵输入为列主序。

编译选项

对于naïve版本的代码，如下所示，不妨先“无脑”地加上-O3 -fomit-frame-pointer -march=armv8-a -ffast-math等编译选项来让编译器尽可能提供些自动向量化的效果。

void square_sgemm (int n, float* A, float* B, float* C) {
   
   
  /* For each row i of A */
  for (int i = 0; i < n; ++i)
    /* For each column j of B */
    for (int j = 0; j < n; ++j)  {
   
   
      /* Compute C(i,j) */
      float cij = C[i+j*n];
      for( int k = 0; k < n; k++ )
        cij += A[i+k*n] * B[k+j*n];
      C[i+j*n] = cij;
  }
}

仅仅是如此，在不同规模的算例上性能就已经有2~10倍的提升，如下图所示。
在这里插入图片描述
可以看到n每逢4的倍数便有显著的性能下降，这是cache thrashing导致的。可做半定量分析：课程集群L1 cache为64B/line，4路组相联，256个组，可知地址低6位为Offset，中间8位为Index，高位为Tag。N-way set associativity只是提供了conflict miss时的“容错性”，因此不失一般性，假定为direct-mapped来分析。地址每隔2^14B就会拥有相同的Index而被映射到同一个set上，对于单精度浮点数而言就是4096个数，因此当n满足(n*m)%4096==0时（m=1,2,…,n-1），就会在一轮k维的循环中产生cache conflict miss，m就是冲突发生时两个B元素相隔的行数。因此冲突频率随n增大而增大，当n≥4096时，就是每两次相邻的对B元素读取都会造成冲突。

循环变换

注意到在naïve的代码中，由于矩阵采用列主序的存储方式，因此先行后列的方式来计算C中元素的值，虽然对B元素访存是连续的，但对于C和A矩阵的访存都是不利的。尤其在循环最内维的k维，A[i+k*n]是大跨步跳跃式访存。
因此可以采用对i和j维的循环交换，来发掘数据复用的空间局部性。代码如下所示。

void square_sgemm (int n, float* A, float* B, float* C) {
   
   
  for (int j = 0; j < n; j++){
   
   
    for (int i = 0; i < n; i++){
   
   
      register float b = B[j*n + i];
      for (int p = 0; p < n; p++)
        C[j*n+p] += A[i*n+p] * b;
    }
  }
}

示意图如下图，相当于按列主序遍历B中元素，对于其中的每个元素b，找到它对应有贡献的C和A中的元素所在的列，进行乘加计算。最内维的p维循环对A和C都是连续的，可以有效利用向量化。由于更改循环后，在整轮最内维的p循环中，b的元素是固定不变的寄存器变量，因此不再出现步骤一中的cache conflict miss，反而是矩阵规模n每逢4的倍数就比相邻的有提升，这是因为n为4的倍数能刚好被向量化指令覆盖，而不会多出额外的数据需要标量运算。
在这里插入图片描述

消除指针别名

消除指针别名告诉编译器修改指针指向的内存内容只能经过该指针之手，使编译器有更大优化空间。主要方法是给函数形参中的指针添加__restrict__关键字。其它局部的指针变量在定义时也可用此修饰。
在这里插入图片描述

循环变换和消除别名的性能均有明显提升，如下图所示。
在这里插入图片描述

循环展开

将循环展开，同时做多列的乘加操作，即取同行不同列的B矩阵元素b0, b1, b2, b3，均与相同的A列做乘法后加到不同的C列上。代码如下所示，需要注意处理余下不足4的列。。

int j, i, p;
for ( j = 0; j < ((n)&(~3)); j+=4)//for each colum j of B
  for ( i = 0; i < n; i++){
   
   //for each row i of B
    register float b0 = B(i,j);
    register float b1 = B(i,j+1);
    register float b2 = B(i,j+2);
    register float b3 = B(i,j+3);
    for ( p = 0; p < n; p++){
   
   
      C(p,j  ) += A(p,i) * b0;
      C(p,j+1) += A(p,i) * b1;
      C(p,j+2) += A(p,i) * b2;
      C(p,j+3) += A(p,i) * b3;
    }
  }
for ( ; j < n; j++)//for each remaining colum j of B
  for ( i = 0; i < n; i++){
   
   //for each row i of B
    register float b0 = B(i,j);
    for ( p = 0; p < n; p++)
      C(p,j  ) += A(p,i  ) * b0;
  }

实验效果显示选4列为一批做乘加效果较好，而大于4列则效果开始下降。循环展开常见的是对最内层做，优势在于循环开销（如终止条件上的分支和计数器变量的更新）的减少。至于为什么要在最外层循环做展开（而不是最内层循环），需要从访存优化的角度来看。对比上一节《循环变换》中最内层循环只有一句C[j*n+p] += A[i*n+p] * b;，展开后此处最内层循环有四句C(p,j ) += A(p,i) * b0;。注意，改写后，A(p,i)只需要载入寄存器一次，就能服务于C(p,j )，C(p,j+1)，C(p,j+2)，C(p,j+3)等的计算；而原来，相同的A[i*n+p]值需要为每个C[j*n+p]加载一次。因此，外层循环的展开将矩阵A元素加载次数减少了nb倍（nb为循环展开的项数，这里是4）。更详细的关于外层循环展开的分析（如中、外层同时展开），可以参见这篇技术资料。当然，按这么分析，outermost loop unrolling这种访存优化对于stencil计算是起不到效果的。

在有的书里，也有类似于这样外层循环展开的写法（只有一层循环，但分成多part，循环体内对每part都做一次计算），分析认为是可以更有效地利用执行单元的流水特性，从而在相同的cycle数内完成更多的计算。

在这里插入图片描述
依然存在n不能整除4时性能下降的情况，因为需要额外执行一段代码。

内存对齐和简单Blocking

利用分块技术提高计算访存比获得更高的性能是常用的优化手段。从之前的代码来看，有三层循环（从外到内依次是j -> i -> p），因此可以在这3个维度上采取分块，分别设为SET_J_BLOCK_SIZE, SET_I_BLOCK_SIZE, SET_P_BLOCK_SIZE。越内维访存越连续，因此设的分块大小更大。此处同时配合内存对齐的手段，是因为对于每一个分块矩阵的乘法，单独将A和B拷贝到一块对齐的连续的内存A_local和B_local中，计算结果存到同样对齐的连续的C_local中。一个好处是A_local和B_local矩阵在拷贝时已经预热，放进了CPU的cache里；另一个好处是在真正计算时，读取和存储都是连续的，提高了cache效率。将一块realMxrealN大小的矩阵拷贝到setMxsetN大小的内存中的代码如下所示。

float C_local[SET_P_BLOCK_SIZE*SET_J_BLOCK_SIZE] __attribute__((aligned(64

最低0.47元/天解锁文章