矩阵乘法&优化方法

jquave

我们定义一个矩阵 为一个大小为m行n列的矩阵，那么两个矩阵的乘法可以定义为:




一个简单的矩阵乘法实现

for (int i = 0; i < M; i++)
  for (int j = 0; j < N; j++)
    for (int k = 0; k < K; k++)
      C[j] += A[k] * B[k][j]
性能：

M, K, N	时间(ms)	GFLOPs
1024, 1024, 1024	1998	1.074

循环重排序（reordering）

调整上述实现中j循环和k循环的顺序，那么我们可以得到
for (int i = 0; i < M; i++)
  for (int k = 0; k < K; k++)
    for (int j = 0; j < N; j++)
      C[j] += A[k] * B[k][j]很明显，循环重排之后，对于A，B，C的空间访问局部性都很得到了保证。

M, K, N	时间(ms)	GFLOPs
1024, 1024, 1024	200	10.69

分块（tiling）

分块可以进一步的提升B C矩阵的空间局部性， 我们把C分成多个title，然后针对每一个title，A和B种对应的行tile和列tile也会切成相应大小的tile进行多个小矩阵乘法，最后加和到C的tile中，当我们把tile的大小限定到合适的范围内时，就可以把整个tile填充到cache内，分块的好处就体现在一个block内的计算小到可以被cache容纳。
for (int i_outer = 0; i_outer < iOuterBound; i_outer++) {
  for (int j_outer = 0; j_outer < jOuterBound; j_outer++) {
    for (int k_outer = 0; k_outer < kOuterBound; k_outer++) {
      for (int i_inner = 0; i_inner < iTile; i_inner++) {
        for (int k_inner = 0; k_inner < kTile; k_inner++) {
          for (int j_inner = 0; j_inner < jTile; j_inner++) {
            C[(i_outer * iTile + i_inner) * N +
              (j_outer * jTile + j_inner)] +=
                A[(i_outer * iTile + i_inner) * K +
                  (k_outer * kTile + k_inner)] *
                B[(k_outer * kTile + k_inner) * N +
                  (j_outer * jTile + j_inner)];
          }
        }
      }
    }
  }
}

M, K, N	时间(ms)	GFLOPs
1024, 1024, 1024	203	10.557

向量化（SIMD）

把最内层的连续的多个乘法计算用一个向量乘法完成，需要用到avx256操作。

M, K, N	时间(ms)	GFLOPs
1024, 1024, 1024	94	22.716

array packing

把分块优化技术中对B的一个块的访问转化为空间上连续的一段数组访问，而不是跳跃式的逐段访问，这中优化技术需要我们用一块中间数组来完成。

M, K, N	时间(ms)	GFLOPs
1024, 1024, 1024	81	26.253

写缓存优化（write caching）

类似的，对C的每个块的计算结果，我们在写回的时候，也是跳跃式的逐段写回，理想的情况应该是一个块的计算结果在内存中是连续的放置的，所以我们需要再开一块write cache内存空间，然后每个block的计算结果直接在write cache上读写，最后计算完一个block之后，整块写回C数组中对应的不同数组段上。

M, K, N	时间(ms)	GFLOPs
1024, 1024, 1024	105	20.384

并行

直接对乘法循环最外层进行unroll，有一定效果，单次乘法的时间下降到了16ms，但离线性扩展还很遥远。

M, K, N	时间(ms)	GFLOPs
1024, 1024, 1024	16	130.178

本文例程:

TVM codegen

参考TVM官方的GEMM优化示例
我在实验机器上面测试了一把，最后使用所有6个优化技术，生成如下的IR：
@main = primfn(A_1: handle, B_1: handle, C_1: handle) -> ()
  attr = {"from_legacy_te_schedule": True, "global_symbol": "main", "tir.noalias": True}
  buffers = {C: Buffer(C_2: Pointer(float32), float32, [1024, 1024], []),
             B: Buffer(B_2: Pointer(float32), float32, [1024, 1024], []),
             A: Buffer(A_2: Pointer(float32), float32, [1024, 1024], [])}
  buffer_map = {A_1: A, B_1: B, C_1: C} {
  allocate(packedB: Pointer(global float32x32), float32x32, [32768]), storage_scope = global {
    for (bigN: int32, 0, 32) "parallel" {
      for (k: int32, 0, 1024) {
        packedB[ramp(((bigN*32768) + (k*32)), 1, 32)] = (float32x32*)B_2[ramp(((k*1024) + (bigN*32)), 1, 32)]
      }
    }
    for (m.outer: int32, 0, 32) "parallel" {
      allocate(C.global: Pointer(global float32), float32, [1024]), storage_scope = global;
      for (n.outer: int32, 0, 32) {
        for (m.c.init: int32, 0, 32) {
          C.global[ramp((m.c.init*32), 1, 32)] = broadcast(0f32, 32)
        }
        for (k.outer: int32, 0, 256) {
          for (m.c: int32, 0, 32) {
            C.global[ramp((m.c*32), 1, 32)] = ((float32x32*)C.global[ramp((m.c*32), 1, 32)] + (broadcast((float32*)A_2[(((m.outer*32768) + (m.c*1024)) + (k.outer*4))], 32)*(float32x32*)packedB[ramp(((n.outer*32768) + (k.outer*128)), 1, 32)]))
            C.global[ramp((m.c*32), 1, 32)] = ((float32x32*)C.global[ramp((m.c*32), 1, 32)] + (broadcast((float32*)A_2[((((m.outer*32768) + (m.c*1024)) + (k.outer*4)) + 1)], 32)*(float32x32*)packedB[ramp((((n.outer*32768) + (k.outer*128)) + 32), 1, 32)]))
            C.global[ramp((m.c*32), 1, 32)] = ((float32x32*)C.global[ramp((m.c*32), 1, 32)] + (broadcast((float32*)A_2[((((m.outer*32768) + (m.c*1024)) + (k.outer*4)) + 2)], 32)*(float32x32*)packedB[ramp((((n.outer*32768) + (k.outer*128)) + 64), 1, 32)]))
            C.global[ramp((m.c*32), 1, 32)] = ((float32x32*)C.global[ramp((m.c*32), 1, 32)] + (broadcast((float32*)A_2[((((m.outer*32768) + (m.c*1024)) + (k.outer*4)) + 3)], 32)*(float32x32*)packedB[ramp((((n.outer*32768) + (k.outer*128)) + 96), 1, 32)]))
          }
        }
        for (m.inner: int32, 0, 32) {
          for (n.inner: int32, 0, 32) {
            C_2[((((m.outer*32768) + (m.inner*1024)) + (n.outer*32)) + n.inner)] = (float32*)C.global[((m.inner*32) + n.inner)]
          }
        }
      }
    }
  }
}对于M=1024, K=1024, N=1024的矩阵乘法，耗时14.036 ms，也即142.857GFlops.
综上，手写的实现达到了TVM优化版本的130.178/142.857 ≈ 91%

JamesB

reordering 部分代码 应该时 C[j] += A[k] * B[k][j] 吧

mastertravels77

是的 已经更正 谢谢