在Tile Language目前的设计理念中,我们提供了三种不同的编程接口,分别面向初学者、开发者和写Kernel的专家这三个等级,我们对Level 1的使用者的定义是不清楚目标硬件架构的开发者(例如不知道GPU的cache一般需要用到tiling,把数据层层cache),这个时候用户编写的程序只是一个简单的计算表达式,不包含任何schedule相关的信息(和tvm的设计一样),这种方法理论上来讲对于不太kernel的人是最舒服的,但是因为很多现存未解决的问题:
1. 从调优的角度来看,从计算到具体的schedule的自动调优时间一般要花比较久的时间,虽然一些工作例如MSRA System Research的Roller,利用硬件感知的白盒搜索空间生成缓解了这一问题, 请看BitBLAS对Roller的复现[Tutorial](https://github.com/microsoft/BitBLAS/blob/main/tutorials/1.fast_and_efficient_codegen.ipynb), 但知道并使用这种方法的人的人也不是很多
1. 从Kernel的表达能力角度来看,目前一些主流的复杂的计算并不能够通过schedule描述出来,例如Flash Attention, 虽然理想上应该可以通过推导 matmul+softmax+matmul 简单计算表达式的一种算子融合版本得到,但是这个问题从写文章的时间点来看仍然很困难。
1. 从社区来看,Schedule的代码虽然看起来还是很优雅的(从写schedule和理解schedule的角度来看, 例如bitblas对于matmul dequantize的[schedule模版](https://github.com/microsoft/BitBLAS/blob/main/bitblas/gpu/matmul_mma_dequantize.py)我个人觉得还是写的很通用的),schedule的魅力在于其从一个最原始不包含任何调度信息的计算表达式,在确保正确性不受影响的情况下一步步变换到最终的高性能调度形式。但是schedule实在太难学习和理解了,即使是一个会写schedule表达式的开发者(这部分玩家已经很少了),想要看明白我写的各种schedule模版,继续扩展也是非常非常困难的。其次,很复杂的计算,例如Flash Attention, 其因为设计本身是要在shared memory上进行算子融合所以计算是无法表达的,其次,即使是要强行写出特别针对Flash attention的多算子fuse版本的schedule模版,schedule本身的代码量可能会比cuda还要长(最后,同样受限于社区,tvm的生态逐渐变得不如triton,一些新feature例如tma这些的支持会有点滞后)
于是在搞bitblas的时候我就觉得这一套有点难受(, 于是觉得需要一个类似triton的东西,但是triton的限制也很大,例如不能显式声明内存,不能显式控制线程的行为等,这一点之后分享的时候再讨论讨论。总之目前bitblas的所有kernel实现都已经换成了Tile Lang,自己用起来非常舒服。
话说回来,我们对Level 2的使用者的定义是知道目标硬件架构的开发者,及在cuda上知道shared memory这个概念,知道在cuda上做tile要先把数据load到一个高速的缓存上再进行计算(有点类似triton的开发模式),本文我们以这种模式为例子介绍一下矩阵乘法Kernel的组成。
最后提一嘴Level 3, thread primitives允许用户完全控制每个线程的行为,写这一部分的代码其实就和写PyCUDA差不多了(但是支持多后端,例如HIP等),于是Level 3就是给完全的expert写的了,但本质上,经过LowerTileOP
和LayoutInference
这两个Pass之后,Level 2的代码也会被Lower到Level 3上。
所以非常值得一提的是,我们的设计中这三种语法可以出现在同一个program中,例如BitBLAS里的量化矩阵乘法的program中对于复杂的反量化部分,我们使用了thread primitives来精心控制每个线程的行为,以及利用ptx来加速精度转换的过程,显式在progam中调用mma 来在一些情况下在寄存器中做反量化等,其他的一些操作,例如数据拷贝和Pipeline仍然使用Level 2的编程方式(T.Pipelined, T.Copy)等,代码参考matmul_dequantize_tensorcore_finegrained.py。