1. 从图到机器码：GPU 代码生成到底在做什么
1.1 为什么编译栈的最后一公里最难走
做深度学习编译栈的人都有一个共识：前端图优化做得再漂亮，如果后端代码生成拉胯，整个栈的性能就是纸上谈兵。TensorPlay Stax 这个编译栈我在几个内部项目里跟过一段时间，前三篇分别聊了图捕获、算子融合和内存规划，到了第四篇，终于要碰最硬的一块骨头——GPU 代码生成与动态形状。
先说清楚这个环节在整个编译栈里的位置。一个典型的编译流程是：前端框架导出计算图，中间层做算子融合和布局推导，然后进入代码生成阶段，把高层的张量操作翻译成 GPU 能执行的 kernel 代码。这个 kernel 可能是 CUDA C、可能是 PTX、也可能是通过某种 IR 再降级到目标指令集。TensorPlay Stax 选择的是先生成一套中间表示（IR），再从 IR 降级到具体后端，这样做的好处是后端可插拔，坏处是每一层降级都可能丢信息。
为什么说这是最后一公里？因为前面所有的优化决策，到了这里都要兑现成实际的线程块划分、共享内存分配、寄存器压力控制。一个融合算子在图层面看起来只占一行，但生成 kernel 的时候可能要展开成几百行代码，涉及几十个参数的计算。更麻烦的是动态形状——当张量的维度在编译期不确定时，你没法提前算好线程数、没法静态分配共享内存，甚至连循环边界都是运行时才知道。
我见过太多项目在这个环节翻车：图优化阶段指标很好看，端到端跑起来却比手写 kernel 慢好几倍。问题往往不在算法，而在代码生成时对硬件特性的理解不到位。所以这一篇我会把 TensorPlay Stax 在 GPU 代码生成上的设计思路、动态形状的处理策略、以及实际踩过的坑都摊开讲。
1.2 本文适合哪些人看
如果你正在做编译器后端、高性能计算、或者深度学习推理引擎的优化，这篇内容应该对你有直接参考价值。如果你只是用框架调模型，可能前几节会有点硬，但关于动态形状的排查技巧和性能调优经验，依然能帮你理解为什么有些模型跑得快、有些跑得慢。
我假设读者对 CUDA 编程模型有基本了解，知道线程块、网格、共享内存这些概念。如果完全不熟，建议先补一下基础，否则后面讲线程映射策略时会比较吃力。另外，文中涉及的具体参数和配置都是基于常见实践给出的参考值，实际项目里需要根据目标硬件做调整。
2. 代码生成的整体架构与选型逻辑
2.1 为什么不用现成的代码生成方案
TensorPlay Stax 在代码生成这块没有直接套用现成的方案，而是自己搭了一套 IR 加多级降级的管线。这个决策在当时是有争议的，因为自己维护一套代码生成器意味着要持续跟进硬件迭代，成本不低。但团队最终拍板自研，核心理由有三个。
第一是动态形状的支持需求。现成的方案大多假设形状静态，遇到动态维度要么回退到解释执行，要么生成大量分支代码，性能损失明显。TensorPlay Stax 的目标场景里有不少变长序列和动态 batch 的负载，这个需求绕不开。
第二是融合算子的代码生成质量。图层面融合出来的复合算子，往往没有现成的 kernel 模板可用，需要根据融合模式动态生成代码。如果依赖外部方案，要么改不动，要么改起来很别扭。
第三是调试和可观测性。自研管线可以在每一级 IR 上插入检查点，出问题的时候能快速定位是哪个降级步骤引入的。用现成方案的话，出了问题只能看最终生成的代码，中间过程是黑盒。
当然，自研的代价也很明显。硬件更新的时候，代码生成器要跟着改；不同厂商的 GPU 架构差异要分别处理；性能调优的很多经验需要自己积累。所以这个选择不是普适的，如果你的场景形状固定、算子种类有限，用现成方案可能更划算。
2.2 三级 IR 的设计与职责划分
TensorPlay Stax 的代码生成管线分三级 IR，每一级解决不同层次的问题。
第一级是张量级 IR，这一层还是以张量为单位，描述的是算子之间的数据流和依赖关系。这一层不做线程映射，只做逻辑上的等价变换，比如把连续的 elementwise 操作合并成一个复合算子。
第二级是循环级 IR，这一层把张量操作展开成嵌套循环，循环边界可以是静态的也可以是动态的。这一层开始引入并行维度的概念，标记哪些循环可以并行、哪些需要串行、哪些可以向量化。循环级 IR 是代码生成的核心，大部分优化决策在这一层完成。
第三级是线程级 IR，这一层把循环映射到具体的线程块和线程，决定共享内存的使用、寄存器的分配、以及同步点的插入。这一层已经非常接近最终的 GPU 代码，降级到 CUDA C 或者 PTX 就是相对机械的翻译工作。
三级 IR 的好处是每一层的职责清晰，优化可以在合适的层次做。比如循环展开在循环级 IR 做，线程块大小调整在线程级 IR 做。坏处是降级过程中可能引入冗余，需要在每一级做清理。
2.3 后端目标的抽象与适配
TensorPlay Stax 支持多个 GPU 后端，包括不同厂商的架构。为了屏蔽差异，代码生成器定义了一套后端抽象接口，主要包括几个方面。
线程模型抽象：不同架构的线程组织方式有差异，有的支持更灵活的线程块形状，有的对线程块大小有限制。抽象接口把这些差异封装起来，上层只需要指定并行粒度，具体映射由后端决定。
内存层次抽象：共享内存的大小、bank 宽度、缓存行为在不同架构上不一样。抽象接口提供查询能力，让上层在生成代码时能根据实际硬件参数做决策。
指令集抽象：向量化指令、原子操作、特殊函数在不同架构上支持程度不同。抽象接口提供能力查询，上层根据查询结果选择可用的指令。
这套抽象不是免费的，每次调用都要经过一层间接。所以在性能敏感的路径上，TensorPlay Stax 会把查询结果缓存起来，避免重复查询。另外抽象层只覆盖差异部分，共性部分直接生成，减少开销。
3. 动态形状的核心挑战与处理策略
3.1 动态形状到底难在哪里
动态形状的难点可以归结为一句话：编译期不知道的信息，运行时要能正确且高效地处理。具体展开，有这么几个层面。
首先是线程映射的不确定性。静态形状下，编译器可以精确算出需要多少个线程块、每个线程处理多少数据。动态形状下，这些都要等到运行时才知道。如果生成代码时假设一个上界，运行时实际值远小于上界，就会浪费大量线程；如果假设一个下界，运行时超出就会出错。
其次是共享内存分配。静态形状下，共享内存的需求是编译期常量，可以直接在 kernel 里声明固定大小的数组。动态形状下，共享内存需求随形状变化，要么用动态共享内存，要么生成多个版本的 kernel 分别处理不同范围。
第三是循环边界和向量化。动态维度导致循环边界不确定，编译器没法做循环展开和向量化，因为这些优化依赖已知的迭代次数。运行时才知道边界的话，只能生成保守的代码，性能损失明显。
第四是算子融合的边界。融合决策往往依赖形状信息，比如判断某个维度是否适合做向量化。动态形状下这些判断做不了，融合策略要么保守要么激进，很难两全。
3.2 形状特化与运行时调度的组合拳
TensorPlay Stax 处理动态形状的核心思路是形状特化加运行时调度。具体做法是：编译期生成多个版本的 kernel，每个版本针对一个形状范围做特化；运行时根据实际形状选择最合适的版本。
形状范围的划分不是随便切的，要结合硬件特性和算子特点。比如对于矩阵乘法，可以按 M 维度的大小分档：小 M 用一个版本，中等 M 用一个版本，大 M 用一个版本。每档的线程块配置、共享内存策略、循环展开因子都可以不同。
运行时调度需要维护一个形状到 kernel 版本的映射表。这个映射表在编译期生成，运行时查表选择。查表本身有开销，所以对于频繁调用的小算子，会把最近使用的形状缓存起来，减少查表次数。
这套机制的关键在于形状范围的划分粒度。切得太细，kernel 版本太多，编译时间和代码体积都会膨胀；切得太粗，特化效果不明显，性能提升有限。TensorPlay Stax 的做法是基于 profiling 数据自动划分，先跑一遍收集各形状区间的性能数据，再根据数据决定切分点。
3.3 动态维度的传播与推导
动态形状在计算图里会传播，一个算子的输出形状依赖输入形状，输入形状又依赖上游算子。TensorPlay Stax 在编译期会做形状推导，能静态确定的维度标记为静态，不能确定的标记为动态，并记录动态维度之间的依赖关系。
这个依赖关系在代码生成时很重要。比如一个 reshape 操作，如果输入形状动态，输出形状也动态，但两者之间有确定的映射关系。代码生成器可以利用这个关系，在运行时根据输入形状直接算出输出形状，而不需要额外的同步或查询。
形状推导的精度直接影响代码生成的质量。推导得越精确，能静态确定的信息越多，生成的代码越激进。TensorPlay Stax 在形状推导上做了不少工作，包括符号化推导、区间分析、以及基于约束的求解。这些技术细节展开讲能写一整篇，这里只提一点：推导结果要保守，宁可标记为动态，也不要错误地标记为静态，否则运行时可能出错。
4. 线程映射与内存优化的实操细节
4.1 线程块大小的选择逻辑
线程块大小的选择直接影响 occupancy 和性能。TensorPlay Stax 在线程级 IR 里会根据几个因素决定线程块大小。
首先是寄存器压力。每个线程用的寄存器越多，能同时驻留的线程块越少，occupancy 越低。代码生成器会估算每个线程的寄存器需求，然后反推合适的线程块大小。估算方法是在循环级 IR 上做寄存器分配模拟，统计活跃变量数量。
其次是共享内存需求。共享内存是线程块级别的资源，线程块越大，每个块能分到的共享内存越少。如果算子需要大量共享内存，线程块就不能太大。
第三是并行度。线程块太小，并行度不够，GPU 的 SM 利用率上不去；线程块太大，调度开销增加，尾效应明显。经验值是在 128 到 512 之间，具体取值要看算子特点。
TensorPlay Stax 的做法是给一个默认值，然后根据上述因素调整。调整逻辑是启发式的，没有理论最优解，需要结合 profiling 数据微调。我实测下来，对于 elementwise 类算子，256 是个比较稳的选择；对于 reduction 类算子，128 或 512 往往更好，取决于 reduction 的维度。
4.2 共享内存的分配与 bank 冲突规避
共享内存用得好能大幅提升性能，用得不好反而成为瓶颈。TensorPlay Stax 在共享内存分配上做了几件事。
首先是动态共享内存的使用。对于形状动态的算子，共享内存需求在编译期不确定，用动态共享内存声明，运行时指定大小。动态共享内存的缺点是编译器没法做静态优化，比如没法把共享内存访问向量化。所以对于能静态确定的部分，还是尽量用静态共享内存。
其次是bank 冲突规避。共享内存按 bank 组织，不同线程访问同一个 bank 的不同地址会产生冲突，导致串行化。TensorPlay Stax 在生成共享内存访问代码时，会检查访问模式，如果发现潜在的 bank 冲突，会插入 padding 或者调整访问顺序。
padding 的做法是在共享内存数组的维度上加一个偏移，让相邻行的起始地址错开，避免同一列的元素落在同一个 bank。这个技巧在矩阵转置等操作里很常用。TensorPlay Stax 会自动判断是否需要 padding，以及 padding 的大小。
4.3 寄存器分配与溢出处理
寄存器是 GPU 上最快的存储，但数量有限。代码生成器要尽量把频繁访问的变量放在寄存器里，同时避免寄存器溢出。
TensorPlay Stax 在循环级 IR 上做寄存器分配，用的是图着色算法。每个变量是一个节点，如果两个变量同时活跃，就连一条边，然后给节点着色，颜色数就是寄存器数。如果颜色数超过可用寄存器数，就要把一些变量溢出到本地内存。
溢出到本地内存的代价很大，因为本地内存实际上在全局内存里，访问延迟高。所以代码生成器会尽量通过循环变换减少活跃变量数量，比如把大循环拆成小循环，减少同时活跃的变量。
我踩过的一个坑是：循环展开因子设得太大，导致活跃变量暴增，寄存器溢出严重，性能反而下降。后来把展开因子从 8 降到 4，性能提升了将近 30%。这个教训是：展开不是越多越好，要结合寄存器压力综合考虑。
5. 代码生成的质量验证与性能调优
5.1 怎么验证生成的代码是对的
代码生成最容易出的问题是正确性。生成的 kernel 可能在某些形状下算错，或者在某些边界条件下出错。TensorPlay Stax 有一套验证机制，分几个层次。
第一层是单元测试。每个算子模板都有对应的测试用例，覆盖常见形状和边界形状。测试方法是拿生成的 kernel 和参考实现对比，逐元素检查误差。
第二层是差分测试。对于动态形状，随机生成一批形状，分别用生成的 kernel 和参考实现跑，对比结果。差分测试能发现形状特化逻辑里的错误。
第三层是模糊测试。构造极端形状，比如维度为 1、维度为质数、维度特别大等，看生成的代码是否还能正确执行。模糊测试主要抓边界条件的 bug。
这三层测试跑下来，大部分正确性问题能提前发现。但测试不能保证百分百正确，所以线上还会加一层运行时检查，对关键算子做结果校验，发现异常及时报警。
5.2 性能调优的常用手段
性能调优是个细活，TensorPlay Stax 在这方面积累了一些常用手段。
循环变换：包括循环展开、循环交换、循环分块。循环展开减少循环开销，循环交换改善内存访问局部性，循环分块提升缓存命中率。这些变换在循环级 IR 上做，变换后的效果通过 profiling 验证。
向量化：把标量操作打包成向量操作，减少指令数。GPU 上常用的向量宽度是 4，也就是一次处理 4 个元素。向量化要求数据对齐，且循环边界能被向量宽度整除。动态形状下，边界可能不整除，需要生成处理尾部的代码。
内存访问优化：包括合并访问、预取、以及使用只读缓存。合并访问要求相邻线程访问相邻内存地址，这样一次内存事务能服务多个线程。预取是提前把数据加载到寄存器或共享内存，隐藏内存延迟。
指令调度：调整指令顺序，让依赖链更短，提高指令级并行。这个优化在生成 PTX 时做，需要了解目标架构的流水线特性。
5.3 常见性能陷阱与排查方法
性能陷阱往往很隐蔽，我整理了几个常见的。
陷阱一：occupancy 过低。表现是 GPU 利用率上不去，SM 大量空闲。排查方法是看 occupancy 计算器，检查寄存器用量和共享内存用量是否超标。解决方法是减少寄存器压力或共享内存需求。
陷阱二：内存带宽瓶颈。表现是计算单元利用率不高，但内存带宽跑满。排查方法是看 profiling 工具的内存吞吐指标。解决方法是优化内存访问模式，提高合并访问比例。
陷阱三：bank 冲突。表现是共享内存访问延迟高。排查方法是看共享内存的冲突统计。解决方法是加 padding 或调整访问顺序。
陷阱四：尾效应。表现是 kernel 执行时间比预期长，尤其是小形状下。原因是线程块数量不是 SM 数量的整数倍，最后一批线程块只有部分 SM 在工作。解决方法是调整线程块大小或网格大小。
陷阱五：动态形状下的分支发散。表现是同一 warp 内的线程走了不同分支，导致串行化。原因是形状特化不够细，不同形状的线程混在一个 warp 里。解决方法是细化形状分档，或者用运行时调度把相同形状的调用聚在一起。
6. 动态形状的实战案例与经验总结
6.1 变长序列处理的代码生成
变长序列是动态形状的典型场景。假设有一个 batch 里包含不同长度的序列，每个序列要过一个 attention 算子。静态形状下，会把所有序列 padding 到最大长度，浪费计算。动态形状下，每个序列按实际长度处理，但代码生成要处理长度不一的问题。
TensorPlay Stax 的做法是生成一个支持动态长度的 kernel，线程块按序列分配，每个线程块处理一个序列。序列长度作为运行时参数传入，循环边界根据这个参数计算。为了避免线程浪费，线程块大小设为一个适中的值，比如 128，长度不足 128 的序列只用部分线程。
这个方案的关键是负载均衡。如果 batch 里序列长度差异很大，长序列的线程块执行时间长，短序列的线程块很快结束，导致 SM 利用率不均。解决方法是把长序列拆成多个块，或者把多个短序列合并到一个块里处理。TensorPlay Stax 用的是后者，通过一个预处理步骤把短序列打包，减少线程块数量。
6.2 动态 batch 下的内存规划
动态 batch 是另一个常见场景。batch 大小在运行时确定，导致中间张量的大小也动态变化。内存规划要处理这个问题。
TensorPlay Stax 的内存规划器会为每个张量分配一个内存池，池的大小按最大可能形状预留。运行时根据实际形状从池里分配。这样做的好处是避免频繁的内存分配释放，坏处是内存占用偏高。
另一种做法是按实际形状动态分配，内存利用率高，但分配开销大。TensorPlay Stax 根据张量的生命周期选择策略：生命周期长的张量用内存池，生命周期短的用动态分配。
内存池的大小估算是个技术活。估大了浪费内存，估小了运行时分配失败。TensorPlay Stax 的做法是基于历史数据统计，取一个较高的分位数作为池大小，同时保留扩容能力。
6.3 我踩过的坑和总结的经验
第一个坑是形状推导过于乐观。有一次形状推导把一个动态维度错误地标记为静态，生成的 kernel 假设这个维度是固定值，运行时遇到不同值就出错。后来加了保守检查，推导不确定的一律标记为动态，问题解决。
第二个坑是kernel 版本爆炸。形状分档切得太细，生成了上百个 kernel 版本，编译时间从几分钟涨到几十分钟，代码体积也膨胀了好几倍。后来重新设计分档策略，把版本数控制在十几个以内，编译时间和性能都回到了合理范围。
第三个坑是运行时调度开销。形状查表和 kernel 选择的开销在小算子上占比很高，导致小形状下性能反而比静态版本差。后来加了缓存，把最近使用的形状和对应的 kernel 版本缓存起来，命中缓存时直接跳过查表，开销降下来了。
第四个坑是共享内存 bank 冲突。一个矩阵转置算子，共享内存访问模式没考虑 bank 冲突，性能只有预期的一半。后来加了 padding，性能翻倍。这个坑让我意识到，共享内存的访问模式必须在代码生成时就考虑，不能等到 profiling 才发现。
总结下来，GPU 代码生成和动态形状处理是个系统工程，涉及编译原理、体系结构、以及大量工程细节。TensorPlay Stax 的方案不一定是最优的，但它在动态形状支持上确实解决了不少实际问题。如果你在做类似的事情，我的建议是：先把静态形状的代码生成做扎实，再逐步引入动态形状支持；形状分档要基于数据，不要拍脑袋；性能调优要持续做，每次硬件或负载变化都可能需要重新调。
最后分享一个小技巧：在代码生成器里加一个开关，可以 dump 出每一级 IR 和最终生成的代码。调试的时候非常有用，能快速定位问题出在哪一级。这个开关在生产环境关掉，避免性能影响。
