【CUDA】TileFusion 最新更新https://zhuanlan.zhihu.com/p/2072480265461567633简介官网TileFusion代码仓库https://github.com/microsoft/TileFusionTileFusion 不是一个独立的 kernel 库而是 FractalTensor 框架中的一个 tile-level GPU kernel 编程与生成后端backend的编程库用于支持 FractalTensor 自动生成融合 GPU kernel。它提供类似tile 分块抽象tile 数据移动 APItile 计算 APItile fusion 机制让用户不用直接写threadblockwarpshared memoryregister fragment而是在 tile 层描述计算。例如传统 CUDAThread | Warp | Block | GridTileFusionKernel | Tile | Warp Group | MMA Instruction | Thread也就是把 CUDA 编程抽象提升了一层。举例说明基础GEMM示例下面是一个使用 TileFusion 编写的简单 GEMM通用矩阵乘法内核示例。完整示例见https://github.com/microsoft/TileFusion/blob/master/examples/101_gemm/01_gemm_global_reg/gemm.hpp。using WarpLayout RowMajor2, 2; 3 // 操作数 A 4 using GlobalA GlobalTileInType, RowMajor128, 256; 5 using IteratorA TileIteratorGlobalA, TileShape128, 32; 6 using RegA RegTileBaseTileRowMajor__half, RowMajor8, 8; 7 using ALoader GlobalToRegLoaderRegA, WarpLayout, kRowReuseCont; 8 9 // 操作数 B 10 using GlobalB GlobalTileInType, ColMajor256, 64; 11 using IteratorB TileIteratorGlobalB, TileShape32, 64; 12 using RegB RegTileBaseTileColMajor__half, ColMajor8, 4; 13 using BLoader GlobalToRegLoaderRegB, WarpLayout, kColReuseCont; 14 15 // 输出 C 16 using GlobalC GlobalTileAccType, RowMajor128, 64; 17 using RegC RegTileBaseTileRowMajorfloat, RowMajor8, 8; 18 using CStorer RegToGlobalStorerGlobalC, RegC, WarpLayout;这个例子中TileFusion 就是在帮助你写一个高性能 CUDA GEMM kernel只不过它把 CUDA 里面复杂的 thread/warp/memory 操作提升成了 Tile 抽象也就是用户直接用 TileFusion C API 写 kernel。展开说明1. 传统 CUDA GEMM 怎么写普通 CUDA 思路GPU Kernel Block | -- Warp | -- Thread Thread: load A load B compute store C用户需要自己决定一个 block 算多少 C一个 warp 算多少矩阵thread 如何搬数据shared memory 怎么布局TensorCore MMA 怎么调用。例如A: [M,K] B: [K,N] C: [M,N] Block 0: load: A[0:128,0:32] B[0:32,0:64] compute: C[0:128,0:64] store: C大量细节。2. 这个例子中TileFusion 做了什么TileFusion 把thread warp block memory tensorcore这些细节封装成Tile TileIterator Loader Storer用户的代码变成Global Memory Tile | Loader ↓ Register Tile | GEMM ↓ Register Tile | Storer ↓ Global Memory回到例子中的代码第一部分定义 A 的 tileusing GlobalA GlobalTileInType, RowMajor128,256;意思我的 A 是A[M,K] 这里定义一个128×256的大tile也就是Global Memory ---------------- | | | 128 x 256 | | A tile | | | ----------------然后using IteratorA TileIteratorGlobalA, TileShape128,32;意思是不要一次处理整个128×256而是切成128×32的小块。所以GlobalA 128×256 切: tile0: 128×32 tile1: 128×32 tile2: 128×32 ...Iterator负责for tile in GlobalA: process(tile)RegTile 是什么这里using RegA RegTile BaseTileRowMajor__half, RowMajor8,8 ;意思是Global memory 中的数据最终进入register每个线程自己的寄存器。GPU三级存储Global Memory | Shared Memory | Register | Tensor Core这里为了简单只用了Global | Register所以GlobalA tile 128×32 Loader ↓ RegA 8×8Loader 干什么看using ALoader GlobalToRegLoaderRegA,...它表示一个Global Memory到Register的数据搬运Global Memory ↓ Register调用loader_a(gAs(k), rA);等价于for threads: load A tile distribute to registers但是用户不用写threadIdx.x threadIdx.y ld.global vector load register mappingTileFusion帮用户生成。GEMM在哪里在这里gemm(rA,rB,acc);这个就是C A×B但是注意这里不是普通CUDA矩阵乘。内部可能映射RegTile ↓ warp layout ↓ TensorCore MMA ↓ mma.sync也就是说代码gemm()最后变成Tensor Core instruction整个kernel执行过程kernelfor(int k0;kIteratorA::sc1;k) { loader_a(); loader_b(); gemm(); } storer_c();展开就是Step 1加载AGlobal Memory A tile 128×32 | Loader ↓ Register A 8×8Step 2加载BGlobal Memory B tile 32×64 | Loader ↓ Register B 8×4Step 3TensorCore计算Register A × Register B ↓ Accumulator Register CStep 4写回Register C | Storer ↓ Global Memory C为什么叫 TileFusion因为它不仅仅做 GEMM。例如普通TransformerGEMM kernel ↓ Bias kernel ↓ Activation kernel ↓ LayerNorm kernel很多kernel launchTileFusion希望Tile0: load GEMM activation store Tile1: load GEMM activation store融合一个kernel完成多个操作