上一篇文章中讲到如何利用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要舒服地多，虽然通过自己瞎理解与实验加上在论坛交流了一下，也算是都摸出来怎么实现，但我相信应该还会有更优雅的写法。
这里记录的是我想从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
一台Windows Server 2022，上面装了两块AMD Radeon™ VII的显卡，操作系统的和显卡的版本都很新，导致出现了一些错误，不过最后还是成功安装上了。
最近在使用NNFusion的时候发现Codegen出来的FP16的网络在V100上的性能打不过FP32（甚至要慢一倍以上），但是理论上FP16应该要比FP32有两倍的性能收益才对（V100 Cuda Core的half precision的最大吞吐量是single的两倍，在s9234的slides中看到直接使用half的情况下peak performance其实和single差不多，都是15Tops，但是Cuda Core提供了half2类型，可一次做两个half类型的运算，这是half在CUDA Core上的收益来源；V100卡上的Tensor Core只支持FP16，利用好Tensor Core可以获得非常强的加速，A100卡上的Tensor Core增加更多的精度支持）。
建议阅读的文章：
拿nvprof测试了一下发现主要的性能瓶颈是：half卷积算子的实现速度要比single慢一倍，而这部分运算又占了总体运行时间的绝大部分。
刚开始碰到的问题是这样的：在Azure上开的一台HPC（4块 V100 16G）在运行了大概七八个小时之后，nvidia的显卡会挂掉，具体的表现为nvidia-smi会卡住十几分钟，之后输出No devices were found，但是执行lspci | grep -i nvidia还是可以看到四块显卡好好的挂在上面，这种情况应该直接reboot就可以修复，但是reboot了之后同样的程序运行一段时间之后显卡还是会掉。
Update your browser to view this website correctly. Update my browser now