0x1. OpenAI Triton介绍阅读,这里来看官方的介绍：https://openai.com/research/triton ，从官方的介绍中我们可以看到OpenAI Triton的产生动机以及它的目标是什么，还可以看到一些经典算法的实现例子展示。
这里的标题是 Introducing Triton: Open-source GPU programming for neural networks ，翻译就是《介绍 Triton：用于神经网络的开源 GPU 编程语言》。然后下面的一句话翻译过来是：我们发布了 Triton 1.0，这是一种开源的类 Python 编程语言，它使得没有 CUDA 经验的研究人员能够编写高效的 GPU 代码——大多数情况下，其效能与专家所能编写的代码相当。这里指出了triton的目的，就是让编写cuda kernrl变得更简单。接下来就逐步看一下介绍里的具体内容，为了更加准确这里会截图对应的原文然后放上我的翻译或者理解。 
这里的意思是Triton可以使得用户用较少的努力就写出一个达到硬件峰值性能的kernel，比如使用 Triton 可以编写 FP16 矩阵乘法的核函数，其性能能够匹配 cuBLAS，并且这个代码不超过25行。然后研究者已经用Triton开发了一些高效的实现，和功能相同的Torch实现相比，性能可以达到两倍提升。后面一段就是强调了使用CUDA来把一些原始的PyTorch实现写一个算子一般会更加高效，但是这个难度不小，并且目前已有工作也不能很好覆盖这种情况，所以OpenAI Triton诞生。 
这里讲的是GPU编程的挑战，现代 GPU 的架构大致可以分为三个主要部分——DRAM、SRAM 和 ALU。在优化 CUDA 代码时，必须考虑到这些组件>
从 DRAM 的内存传输必须合并成大型事务，以利用现代内存接口的大总线宽度（内存合并访问）。
数据必须在重复使用前手动存储到 SRAM 中，并进行管理来最小化bank conflict。
计算必须仔细地进行划分和调度，不仅是在流式多处理器（SMs）之间，还包括在其内部，以促进指令/线程级并行性，并利用专用的 ALU（例如，Tensor Cores）。 
考虑所有这些因素可能对于拥有多年经验的资深 CUDA 程序员来说都是一个挑战。Triton 的目的是完全自动化这些优化，以便开发者能够更好地专注于他们并行代码的高层逻辑。Triton 旨在广泛适用，因此不会自动在流式多处理器（SMs）之间调度工作——留下一些重要的算法考虑（例如，tiling，跨 SM 同步）由开发者自行决定。
然后给了一个表格展示cuda的编译器和triton的区别。 
在所有可用的领域特定语言和即时编译器中，Triton可能和Numba最相似：kernel被定义为一个装饰过的函数，并以不同的 program_id 并行启动在所谓的网格实例上。然而，正如下面的代码片段所示，相似之处仅此而已：Triton 通过对块上的操作来暴露实例内部的并行性——这些小数组的尺寸是二的幂次方——而不是单指令多线程（SIMT）执行模型。这样做，Triton 有效地抽象出了所有与 CUDA 线程块内部并发相关的问题（例如，内存合并、共享内存同步/冲突、Tensor Cores调度）。 
注意，Triton 的即时编译器将 X 和 Y 视为指针而不是张量；我们认为保留对内存访问的低级控制对于处理更复杂的数据结构（例如，块稀疏张量）是重要的。重要的是，这种特定的 softmax 实现在整个标准化过程中将 X 的行保留在 SRAM 中，这在适用时最大化了数据重用（约 <32K 列）。这与 PyTorch 的内部 CUDA 代码不同，后者使用临时内存使其更具通用性，但显著更慢（如下所示）。这里的关键不是 Triton 本质上更好，而是它简化了专用kernel的开发，这些内核可能比在通用库中找到的内核快得多。 
Torch（v1.9）JIT编译器的较低性能凸显了从高级张量操作序列自动生成 CUDA 代码的难度。 
这里是说Triton大概只需要25行Python代码就可以实现一个接近峰值的矩阵乘法。（后面有专门的一大节讲这个代码的原理）代码如下>
手写矩阵乘法kernel的一个重要优势是，它们可以根据需要定制，以适应输入（例如，切片）和输出（例如，LeakyReLU）的融合转换。如果没有像 Triton 这样的系统，没有出色的 GPU 编程专长的开发者将无法进行矩阵乘法内核的定制修改。 
这里是说Triton 的良好性能源于一个以 Triton-IR 为中心的模块化系统架构，Triton-IR 是一个基于 LLVM 的中间表示，在这个系统中，多维值块（这个是MLIR的概念）是一等公民。GPT,@triton.jit 装饰器的工作原理是遍历提供的 Python 函数的抽象语法树（AST），以便使用常见的 SSA 构建算法即时生成 Triton-IR。然后，编译器后端会简化、优化并自动并行化所产生的 IR 代码，再将其转换为高质量的 LLVM-IR —— 最终生成 PTX —— 以在近期的 NVIDIA GPU 上执行。目前不支持 CPU 和 AMD GPU，但我们欢迎社区贡献，旨在解决这一限制。 
我们发现，通过 Triton-IR 使用块级别程序表示，使我们的编译器能够自动执行各种重要的程序优化。例如，可以通过观察计算密集型块级操作（例如，tl.dot）的操作数，自动将数据暂存到共享内存中，并使用标准的活性分析技术进行分配和同步。
另一方面，如下所示，Triton 程序可以高效且自动地并行化，既可以（1）通过并发执行不同的kernel实例在流式多处理器（SMs）间并行，也可以（2）通过分析每个块级操作的迭代空间，并在不同的 SIMD 单元间适当分配，从而在 SMs 内部并行。 
0x2. 教程1 Vector Addition阅读,
意思是这一节教程会介绍Triton编程模型定义kernel的基本写法，此外也会介绍一下怎么实现一个良好的benchmark测试。下面来看计算kernel实现，我把注释改成中文了>
这里还声明了一个辅助函数来（1）分配z张量，（2）使用适当的网格/块大小排队上面的kernel>
我们现在可以使用上面定义的函数来计算两个torch.tensor对象的逐元素求和，并测试其正确性>
输出>
我们可以对不同大小的向量进行自定义操作的性能基准测试，以了解它相对于PyTorch的表现如何。为了简化操作，Triton提供了一系列内置工具，使我们能够简洁地绘制出自定义操作在不同问题规模下的性能图表。
gbps = lambda ms: 12 * size / ms * 1e-6这里的12表示的是数据读写的bit，因为有x和y以及z的存在，所以是3*4=12bit。现在可以运行上面的装饰函数了。传递 print_data=True 参数来查看性能数据，传递 show_plots=True 参数来绘制图表，和/或传递 save_path='/path/to/results/' 参数来将它们连同原始CSV数据一起保存到磁盘上> 
可以看到，对于elementwise任务，Triton的性能几乎和PyTorch持平，但是Triton写起来很简单。
0x3. 教程2 Fused Softmax阅读,在这个教程中，我们将编写一个融合的softmax操作，这个操作对于特定类型的矩阵来说比PyTorch的原生操作要快得多：那些行的大小可以放入GPU的SRAM中的矩阵。
通过这样做，我们将学习到>
kernel融合对于带宽受限操作的好处。
Triton中的reduce操作符。
动机,自定义GPU kernel用于逐元素加法在教育上是有价值的，但在实际应用中可能作用有限。让我们考虑一个简单的（数值稳定的）softmax操作的情况> 
计算kernel,我们的softmax kernel的工作方式如下：每个程序加载输入矩阵X的一行，对其进行归一化处理，然后将结果写回到输出Y中。需要注意的是，Triton的一个重要限制是每个块必须包含2的幂次方个元素，因此如果我们想处理任何可能的输入形状，我们需要在内部对每行进行“pad”以及对内存访问操作进行保护（也就是防止越界）>
解析来创建一个辅助函数，该函数为任何给定的输入张量排队执行kernel并且设置了启动参数。 
这里是验证Triton实现的fuse softmax和PyTorch的naive实现等价，显然他们是等价的。
BenchMark,
这里设定矩阵的行数为固定的4096来做benchmark。 
这里提到虽然Triton实现的softmax性能更好并且易于理解和维护，但PyTorch的torch.softmax则更加通用。
0x4. 教程3 Matrix Multiply阅读,
首先教程指出这里就是要写一个Block级别的矩阵乘法，然后这里会涉及到多维度的指针操作，程序重排以更好的命中l2 cache以及自动调优。
动机,矩阵乘法是大多数现代高性能计算系统的关键构建块。它们众所周知难以优化，因此它们的实现通常由硬件供应商自己作为所谓的“内核库”（例如，cuBLAS）的一部分来完成。不幸的是，这些库通常是专有的，无法轻易地定制以适应现代深度学习工作负载的需求（例如，融合激活函数）。在这个教程中，你将学习如何使用Triton自己实现高效的矩阵乘法，这种方法易于定制和扩展。
大致来说，我们将要编写的内核将实现以下块级算法来乘以一个 (M, K) 矩阵和一个 (K, N) 矩阵>
其中，双重嵌套的for循环的每次迭代都由一个专用的Triton program实例执行。
计算kernel,上述算法实际上在Triton中相当容易实现。主要的难点来自于在内循环中计算必须读取A和B块的内存位置。为此，我们需要多维指针运算。
指针运算,对于一个2D Tensor X，X[i, j]的内存位置为&X[i, j] = X + i*stride_xi + j*stride_xj。因此，对于A[m : m+BLOCK_SIZE_M, k:k+BLOCK_SIZE_K]和B[k : k+BLOCK_SIZE_K, n : n+BLOCK_SIZE_N]的块指针可以用下面的伪代码定义>
这意味着A和B块的指针可以在Triton中初始化，比如 k=0 如下代码所示。另外注意，我们需要一个额外的模运算来处理M不是BLOCK_SIZE_M的倍数或N不是BLOCK_SIZE_N的倍数的情况，在这种情况下，我们可以用一些无用的值填充数据，这些值不会对结果产生影响。对于K维度，我们稍后将使用掩码加载语义来处理。
然后在内循环中按如下方式更新>
如上所述，每个program实例计算一个 [BLOCK_SIZE_M, BLOCK_SIZE_N] 大小的C矩阵块。重要的是要记住，这些块的计算顺序是很重要的，因为它会影响我们程序的L2缓存命中率，不幸的是，一个简单的行优先顺序是不够的。
L2 Cache优化,如上所述，每个程序实例计算一个 [BLOCK_SIZE_M, BLOCK_SIZE_N] 大小的C矩阵块。重要的是要记住，这些块的计算顺序很重要，因为它会影响我们程序的L2缓存命中率，不幸的是，一个简单的行主序排序是不够的。
一个可能的解决方案是以一种促进数据重用的顺序启动块。这可以通过在切换到下一列之前将块在GROUP_M行的super group中分组来实现>
例如，在下面的矩阵乘法中，每个矩阵由9个块乘以9个块组成，我们可以看到，如果我们按行主序计算输出，我们需要将90个块加载到SRAM中以计算前9个输出块，但如果我们按grouped ordering进行计算，我们只需要加载54个块。 
在实际应用中，这可以在某些硬件架构上提高我们矩阵乘法内核的性能超过10%（例如，在A100上从220提升到245 TFLOPS）。
L2 Cache优化原理补充讲解,上面的group oredering的访问代码比较难理解，这里来更详细的解析一下。 
这里的num_pid_m和num_pid_n就是求分别要在M和N方向循环多少次。
然后上面图中的黑色数字其实就可以理解为program id，我们可以看到program id增加的方向其实就代表了遍历的ordering，对于row major来说就是在行方向上顺序遍历，而对于group ordering来说就是按照一个BLOCK_SIZE_M*BLOCK_SIZE_N这么大的一个小组来遍历。其实这段代码就是完成group ordering的遍历>
以上面图来看，num_pid_m=3，num_pid_n=3，num_pid_in_group=group_id * GROUP_SIZE_M=9*3=27，也就是下面的红色框里面的program个数，从名字也可以看出来这个红色框划分的区域也是一个group。 
group_id 就表示当前的这次 "循环", 是在第几个红色框里，以program 0为例，这里为group_id = pid // num_pid_in_group=0//27=0。而first_pid_m 代表当前 group 中的第一个黄色program在全局的M维度上是第几个program ，这里为first_pid_m = group_id * GROUP_SIZE_M=0，group_size_m = min(num_pid_m - first_pid_m, GROUP_SIZE_M)这里是考虑到最后一个group可能占不满数据（存在padding），所以就做一个截断处理。
这两行代码计算当前的program处理的黄色小块坐标（[pid_m, pid_n]），pid_m这行是在行方向上移动，pid_n这行则是保证在上面的红色框里面一定是一列一列来访问的。
作为对比，在Row-major的方法中，访问方式应该是这样的>
计算最后的结果,有了上面的铺垫，我们就可以计算最终的结果了，下面的代码展示了完整的Triton 矩阵乘法kernel实现。
我们现在可以创建一个方便的封装函数，它只需要两个输入张量，并且会：（1）检查任何形状约束；（2）分配输出；（3）启动上述kernel。
计算过程的补充说明,上面的《L2 Cache优化原理补充讲解》这一节明确了kernel的group ordering的访问方式以及实现，现在来看对于当前的program实例具体是怎么计算的。现在以计算C中的第一个Block的(0, 0)为例子，它需要从A和B分别加载9个黄色的小块数据相乘并累加最后得到C中的(0, 0)位置结果。如下图所示> 
下面的代码先把program实例当前要处理A和B的第一个Block加载上来>
这里的a_ptr 是整个 A 矩阵第一个元素的地址，offs_am和offs_bn表示当前的program id在M维度和K维度的坐标，这个坐标是一个list，用tl.arange(0, BLOCK_SIZE_K)来获取。
得到 M 维度 和 K 维度的坐标后, 就可以让它们各自和 M 维度 和 K 维度的 stride 相乘, 然后和 a_ptr 相加, 就可以得到 A 矩阵 9 个 block 中第一个 block 中每个元素的地址了。 b_ptr也是同理。
最后一部分就是累加了，这里会在K维度上进行累加，每次计算输出的一个块。
这行代码a = tl.load(a_ptrs, mask=offs_k[None, :] < K - k * BLOCK_SIZE_K, other=0.0)考虑到 K 可能不能被 BLOCK_SIZE_K 整除, 到每一行最后一个 block 的时候, 实际大小是不足 BLOCK_SIZE_K 的，所以需要把超出的那部分元素mask掉。
最后这部分代码是把当前的算子和LeakyReLU激活函数进行融合>
单元测试,
Benchmark,这里使用一个方阵来对比Triton实现的matmul kernel和cublas的matmul kernel的性能。
可以看到基于Triton实现的矩阵乘kernel性能大体可以和高度优化的cuBlas持平。
