在之前的两篇文章中，我们分别用TVM的Tensor Expression与TIR Script完成了在Nvidia Cuda Core上的高效的FP32 矩阵乘法，3090-24GB的各种精度在Cuda Core和Tensor Core上的Peak TFLOPS如下表所示：

有意思的是，3090上，FP16的Peak Peformance和FP32是一样的，这一点比较特殊，是因为架构上的改动，一般而言fp16的性能都会是fp32的两倍或者四倍，这个主要是因为20系的gpu把fp32和int32的Cuda Core分开了，从而能同时进行fp32和int32的计算，30系把int32的core又就加上了fp32的计算单元，所以fp32的计算能力翻倍，而cutlass下的16384的gemm。

按照3090上的硬件单元分类，我们还可以探索一些有意思的加速，比如在CUDA Core上使用SIMD指令（DP4A，HFMA2来优化int8、half的性能，

Read More

上一篇文章中讲到如何利用cutlass优化gemm的思路，使用tvm tensor expression来实现一个高效的矩阵乘法，这里再探索一下直接从TIR Script把这个东西复现一下，对比一下两者的异同。

TensorIR 今年7月再arxiv上放了一篇preprint，感兴趣的读者可以自行阅读：https://arxiv.org/abs/2207.04296

不过写这篇文章的时候，tvm上游(main)分支的tir与paper里还不是一样，siyuan他们另做了许多改进，估计要等paper中了才会被合并到上游（貌似是在投ASPLOS？所以这里还是以tvm上我们可以实际操作的TensorIR Script为例子，优化的思路则不多讲解，和之前的tensor expression是一样的。

PS: 感觉TIR Script的设计和写法更贴近GPU，比tensor expression更抽象，有亿点点摸不着头脑，不过也比直接从tensor ir来构建一个dag要舒服地多，虽然通过自己瞎理解与实验加上在论坛交流了一下，也算是都摸出来怎么实现，但我相信应该还会有更优雅的写法。

代码还是放在: https://github.com/LeiWang1999/tvm_gpu_gemm

Read More

这里记录的是我想从tvm的tensor expression出发，参考一下cutlass efficient gemm的思路，一步一步优化一下GEMM的一些思考和碎碎念，目的是为了理解cutlass优化gemm的思路。

我们使用CUTLASS Profiler来运行一个gemm的运算，并用nsight compute dump下来其运行过程中的一些情况，可以拿到他的一些信息，如grid的大小与block的大小等。比如对于16384的float32类型数据的gemm，cutlass的grid size是（512, 16, 1）-> 8192个block， block size是（256，1，1），一共是2,097,152个线程，因为最后产生C的大小是（16384，16384），所以平均每个thread需要产生128个C的元素，结合这些参数的信息，使用tvm的te进行schedule（其实可以试试tensor ir），最后成功打到了和cublas，cutlass相近的性能。

测试GPU: rtx 3090 24GB

CUDA Version: 11.1

TVM Version: 10.0

代码放在：https://github.com/LeiWang1999/tvm_gpu_gemm

Read More