上一篇文章中讲到如何利用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

Digilal DesignEEEE

这里记录的是我想从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

CUDA ProgrammingMLSys

最近在使用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增加更多的精度支持)。

建议阅读的文章:

聊聊 GPU 峰值计算能力

A100 Tensor Float 32 性能实测

拿nvprof测试了一下发现主要的性能瓶颈是:half卷积算子的实现速度要比single慢一倍,而这部分运算又占了总体运行时间的绝大部分。

Digilal DesignEEEE