08_turing_tensorop_gemm

在00_basic_gemm中基本梳理清楚了gemm计算流程，不过使用cuda core写的，虽然目前没人用这个实现gemm了，但是优化思想还是值得学习的，现在来看看turing架构下的gemm怎么实现，和basic gemm的区别就是mma是tensor core, 具体如下：

和之前相比，区别在于做warp gemm的时候，上面的小绿色块是整个warp去计算得到的，之前是一个线程做。因此整个warp tile需要循环做8*8次才可以，具体实现在mma_tensor_op.h,代码如下：

      CUTLASS_PRAGMA_UNROLL
      for (int m = 0; m < MmaIterations::kRow; ++m) {

        CUTLASS_PRAGMA_UNROLL
        for (int n = 0; n < MmaIterations::kColumn; ++n) {

          int n_serpentine = ((m % 2) ? (MmaIterations::kColumn - 1 - n) : n);

          if (AccumulatorsInRowMajor) {  // matrix B is reordered
            mma(
              ptr_D[n_serpentine + m * MmaIterations::kColumn],
              ptr_A[m],
              ptr_B[n_serpentine],
              ptr_D[n_serpentine + m * MmaIterations::kColumn]);
          } 
        }
      }
    }

对于tensor core的具体使用，需要配合文档才能使用，因为这里的mma是warp为概念的，对于每个线程只要做搬运数据搬运工作就行，程序员的作用就是根据文档，让每个线程搬运指定位置的数据到寄存器就行，得到的数据再放到指定位置就行，nv这块文档给的不详细，后来推荐用wmma的api，想搞细节可以看看tensor core, 不想搞的话用用wmma也行。tensor core的实现在mma_sm75.h

  CUTLASS_HOST_DEVICE
  void operator()(
    FragmentC &d,
    FragmentA const &a,
    FragmentB const &b,
    FragmentC const &c
  ) const {

  unsigned const & A = reinterpret_cast<unsigned const &>(a);
  unsigned const & B = reinterpret_cast<unsigned const &>(b);

  int const *C = reinterpret_cast<int const *>(&c);
  int *D = reinterpret_cast<int *>(&d);

  asm volatile("mma.sync.aligned.m8n8k16.row.col.satfinite.s32.s8.s8.s32 {%0,%1}, {%2}, {%3}, {%4,%5};\n"
      : "=r"(D[0]), "=r"(D[1])
      : "r"(A), "r"(B), "r"(C[0]), "r"(C[1]));
  }
};

原文地址：https://blog.csdn.net/feng__shuai/article/details/145201043

免责声明：本站文章内容转载自网络资源，如侵犯了原著者的合法权益，可联系本站删除。更多内容请关注自学内容网（zxcms.com）！