最近在调torch.compile生成的 kernel 时发现一个很值得琢磨的现象同样一个 pointwise 表达式在 CPU 后端会生成带#pragma omp和 SIMD 循环的 C 代码在 GPU 后端则会生成一段triton.jit的 Python 代码。这两条路径在 TorchInductor 内部就是由CppKernel和TritonKernel两个内核生成器完成的。之前系列第一篇聊了从 FX Graph 到 Inductor IR 的整体流程这篇就专门拆开CppKernel与TritonKernel看看同一份 IR 是怎么分别落到两种完全不同的硬件代码上的以及实际调试时会碰上哪些坑。1. 两条生成路径为什么长成了两副面孔1.1 CPU与GPU的硬件约束完全不同看代码生成器先看硬件。CPU 核心少但单核强执行模型是 多线程 SIMD编译器帮你把循环展开成 AVX/SSE 指令但前提是循环边界清楚、内存连续、指针不对齐的时候还得兜底处理。GPU 核心多但单线程弱执行模型是 大量轻量线程 块级调度你如果给 GPU 也生成一层层 C 风格的标量循环那基本就是在用一万人排队过独木桥完全浪费吞吐量。所以 TorchInductor 从一开始就不能只维护一套内核生成逻辑。CPU 上要生成的是面向循环优化的 CGPU 上要生成的是面向块级并行的 Triton。这两者虽然都叫 kernel内部结构、执行模型、优化思路完全不是一个物种。理解这一点后面看CppKernel和TritonKernel的代码分工才有根基。1.2 CppKernel与TritonKernel的类职责划分在 TorchInductor 源码里CppKernel对应torchinductor/codegen/cpp.pyTritonKernel对应torchinductor/codegen/triton.py。两个类要解决的问题是一致的拿到 Inductor IR 中排好序的SchedulerNode生成一个可执行的 kernel 函数体。但由于目标设备不同各自需要处理的细节完全不同。CppKernel需要关心的事情包括C 函数签名怎么声明、临时 buffer 怎么分配、循环嵌套怎么展开、尾数循环怎么处理、OpenMP 并行度怎么设、SIMD 向量宽度怎么确定TritonKernel关心的是Triton 函数签名怎么写、tl.program_id怎么映射、BLOCK 大小怎么设、mask 是否需要、reduction 用tl.sum还是原子操作。它们还有一个共同的上游组件叫Scheduler。Scheduler负责把 Inductor IR 根据依赖关系切分成一个个 kernel 边界然后根据设备类型选择一个生成器。你可以把Scheduler想成装修公司的项目经理它只负责哪里该砌墙、哪里该布线但具体到墙怎么抹灰、线怎么穿管是泥瓦队和电工队各自的事。1.3 共用上游IR但生成逻辑彼此独立这里的关键点是Inductor IR 本身是设备无关的。Pointwise、Reduction、Scan这些节点只表达计算逻辑不绑定 CPU 还是 GPU。CppKernel和TritonKernel都是把这个 IR翻译成各自的编程语言而不是各自发明一种新 IR。这个设计带来的直接好处是你在高层做算子融合、做 buffer 复用、做依赖分析时不用关心后端差异。一旦进入生成阶段两边就各走各路。举个例子同一个Pointwise节点在CppKernel里会被映射成一层for循环循环体里是float或向量Vec类型的运算在TritonKernel里会被映射成一段tl.load 计算 tl.store并行单位从循环迭代变成了block 内的一整块连续数据。所以如果你只想读懂 TorchInductor 的代码生成建议先同时打开cpp.py和triton.py对照同一个 IR 节点看两边各自怎么处理。比只看单边更容易理解哪些是设备无关的哪些是硬件逼出来的。2. CppKernel把 Inductor IR 压实成能跑AVX的C2.1 LoopNest、Buffer和kernel签名从哪来CppKernel生成的并不是一个独立可编译的 C 文件而是一段以extern C void kernel(...)形式出现的函数体。函数参数来自 scheduler 为这个 kernel 分配的输入输出 buffer以及一些必要的标量参数。比如y x 1这样的算子最终生成的签名大致是void kernel(float* x_ptr, float* y_ptr, long n)。这里的核心数据结构是LoopNest。Inductor IR 里的Pointwise节点在 C 侧会被展开成一个多层嵌套循环每一层循环对应一个维度。CppKernel 需要计算每个循环维度的起始、结束和步长还需要决定哪些循环可以内层连续访问、哪些循环适合放在外层交给 OpenMP 并行。这些决策直接影响生成的循环性能。还有一个很容易忽略的细节是 buffer 对齐。CPU 向量化加载要求数据地址满足一定对齐条件比如 AVX2 下 float 数组要 32 字节对齐。TorchInductor 在分配 buffer 时会尽量满足这个条件但在经过视图、切片、非连续张量后对齐信息可能丢失。CppKernel 在生成加载代码前会检查指针的对齐属性如果无法确认对齐就不能放心生成向量化版本。2.2 向量化路径CppVecKernel与标量回退CppKernel并不是一个最终的生成实现它下面还有一个重要的子类叫CppVecKernel。从名字就能看出这是专门做向量化代码生成的版本。CppVecKernel会把循环步长从1改成VecSize并将循环体内的标量运算替换成向量运算。这个过程很像给编译器递小抄。虽然 GCC 和 Clang 本身都有自动向量化能力但自动向量化依赖非常苛刻的循环条件边界必须是常数、内存访问必须是严格连续、不能有潜在别名。Inductor IR 经过各种变换后循环边界经常是运行时变量这时候保守的编译器会放弃向量化。CppVecKernel的做法是显式生成 一次处理 N 个元素 的循环通过内部的load/store封装把数据加载成向量类型计算完成后存回内存。这里必须提一个常见回退当数据指针无法确认对齐、或者循环尾部剩余不足一个VecSize时CppKernel 会生成一段标量 tail 循环。我在实际调参时经常发现生成的 C 里既有一大段向量化主循环、又有一个for (long i n - n % 8; i n; i)的尾巴。这个尾巴看起来不起眼但如果 n 很小整个 kernel 可能完全退化为标量版本。所以如果你的算子输入 shape 都不足向量宽度那谈 AVX 优化就是空谈。2.3 OpenMP并行与reduction的循环策略CppKernel 在 CPU 上要利用多核通常会在外层循环上加#pragma omp parallel for。这个决策不是无脑加的需要平衡并行开销和内存带宽。比如一个单次遍历内存带宽有限的 elementwise 算子开太多线程反而可能因为争抢内存控制器导致性能下降。reduction 是另一个值得细看的策略。CPU 上的 reduction 如果直接在外层加 omp每个线程会各算一部分偏和最后再做一次跨线程合并。TorchInductor 在 IR 里会告诉CppKernel哪些节点是 reduction哪些是 pointwise。CppKernel 需要决定对 reduction 维度使用局部变量累积还是使用多级循环。这里容易犯的错是忽略 OpenMP 默认共享变量的问题必须显式声明private的累积变量否则多线程写同一个变量就是数据竞争。从源码角度CppKernel的这几件事不是分散在各处而是集中在codegen函数里通过大量条件分支判断当前节点类型、循环层级和并行策略。如果你要改 CPU 生成逻辑建议从CppKernel.codegen入手顺着它对每个SchedulerNode的处理往下追。3. TritonKernel在GPU上把IR变成块级Triton程序3.1 TritonKernel的代码骨架与program_id映射GPU 侧的TritonKernel生成的东西是一段 Python 风格的 Triton 代码。Triton 是一种类似 CUDA 但更高级的编程语言它让你不用手动管理 block 和 warp而是用程序加上块内布局的抽象。TorchInductor 的TritonKernel本质上是在帮你自动写这段 Triton 代码。看一个典型骨架首先生成triton.jit注解的函数函数参数是所有输入输出 buffer 的指针以及必要的标量。函数内部先通过tl.program_id(0)拿到当前程序块的 id然后计算这个 block 对应的起始 offset通过tl.arange(0, BLOCK)得到块内连续的 index 序列再用offset pid * BLOCK arange得到全局面索引。之后就是根据 IR 的计算逻辑生成tl.load、计算、tl.store。这个骨架和 CUDA kernel 的网格-块映射其实是一一对应的。grid 的维度由cdiv(xnumel, BLOCK)决定每个 program 处理BLOCK个元素。TritonKernel要做的就是把这个映射关系算清楚并且把 IR 里的逐元素计算改写成块内向量化的计算。因为 Triton 编译器会帮你把tl.arange映射到线程束所以最终执行效率很大程度上依赖BLOCK选得是否合理。3.2 BLOCK大小、tl.constexpr与自动调优BLOCK是 TritonKernel 里最重要的配置参数之一。它必须是 2 的幂而且不能太小太小则单个程序块占用硬件资源不足也不能太大太大会导致寄存器溢出。TorchInductor 给出的默认策略是在一组候选值中选择比如 128、256、512、1024然后分别生成代码、在真实输入上跑 benchmark选出最快的一个。这个过程里tl.constexpr起了关键作用。Inductor 在生成代码时会把BLOCK、以及一些诸如当前维度是否整除的布尔值标记成tl.constexpr。这样 Triton 编译器在编译时就能把 constexpr 常量直接折叠进代码消除掉所有动态判断。比如mask逻辑如果被 constexpr 标记为不需要生成的代码里就不会出现offs xnumel的比较。我调试时发现一个隐藏含义由于自动调优需要对同一份 IR 生成多份不同BLOCK的代码TritonKernel每次生成过程都必须可重复、无副作用。如果你的自定义生成逻辑里引入了随机顺序依赖调优结果的置信度会大打折扣。3.3 mask、规约和原子操作的生成决策TritonKernel 对非整除边界的处理比 CppKernel 更优雅。CppKernel 需要生成显式 tail 循环而 TritonKernel 通常只是生成一个 maskmask offsets xnumel然后在tl.load和tl.store时传入这个 mask。这样尾部越界的线程不会访问非法内存。但 mask 也不是无脑生成的。如果 Inductor 在 IR 阶段就能证明xnumel % BLOCK 0那TritonKernel就不会生成mask相关的代码因为不需要。这就要求 IR 的 shape 推导信息足够准确。我之前遇到过因为 shape 推导未能识别常量维度导致代码多出 mask 和other0性能下降不少。这类问题往往需要回到 IR 生成侧去看 shape 信息是否在某个 pass 中丢失了。reduction 在 TritonKernel 里的生成也有讲究。最简单的归约比如求和会生成tl.sum(acc, axis0)。但如果归约跨多个 block每个 block 只得到部分结果最后需要对部分结果做合并。TorchInductor 根据 kernel 的归约范围决定是生成单 block 内归约、还是生成跨 block 归约。跨 block 归约通常需要多个 kernel 完成或者使用原子操作。TritonKernel会在 IR 的依赖关系中插入一个中间 buffer让第一个 kernel 把 partial result 写入第二个 kernel 再加载并最终归约。这种多级归约的生成逻辑是 GPU 后端代码生成里最容易出错的部分也是和 CPU 后端差异最大的部分。CPU 上一个简单循环累加就完事GPU 上却要考虑 block 间同步和原子竞争导致代码结构完全不同。4. 同一个表达式两种kernel产物对比4.1 生成C代码的典型形态为了直观展示差异我用一个非常简单的y x 1来举例。经过 TorchInductor 调度后CppKernel 生成的代码骨架大致如下。下面代码是示意风格真实生成代码会包含更多宏和辅助结构extern C void kernel(float* x, float* y, long n) { const long vec_size 8; // 假设是AVX2下的float宽度 long n_vec n - (n % vec_size); #pragma omp parallel for for (long i 0; i n_vec; i vec_size) { Vecfloat xv load_float8(x i); Vecfloat yv xv 1.0f; store_float8(y i, yv); } for (long i n_vec; i n; i) { y[i] x[i] 1.0f; } }可以看到CppKernel 的思路是一个 CPU 线程顺序处理多个元素尽量让单个线程的计算向量化。整个代码的执行模型是扁平的外层 omp 负责跨核内层向量化负责跨数据。内存访问模式是直接指针运算没有显式的 block 概念。4.2 生成Triton代码的典型形态同一表达式的 TritonKernel 产物会是这样triton.jit def kernel(x_ptr, y_ptr, n, BLOCK: tl.constexpr): pid tl.program_id(0) offsets pid * BLOCK tl.arange(0, BLOCK) mask offsets n x tl.load(x_ptr offsets, maskmask) tl.store(y_ptr offsets, x 1.0, maskmask)Triton 代码的执行模型是整个 block 一起处理多个元素。tl.arange(0, BLOCK)在 Triton 编译后会映射到某个 warp 的 lane程序块自动展开成硬件线程束。CppKernel 的循环迭代是显式的TritonKernel 的循环迭代被抽象掉了你写的是块级操作具体线程调度由 Triton 编译器完成。4.3 关键差异对照表把两种产物放在一起看差异非常清晰对比维度CppKernelTritonKernel目标设备CPUGPU编程语言CTriton / Python并行单位OpenMP 线程Triton program / block数据遍历方式显式 for 循环 SIMDtl.arange tl.load/tl.store尾部处理标量 tail 循环mask 掩码向量化显式 Vec 类型块内自动向量化调优方式编译器参数、对齐、omp策略BLOCK 大小、tl.constexpr 特化reduction局部变量 omp reducetl.sum 原子操作/多kernel合并这张表也是一个很好的自查清单。如果你在修改 TorchInductor 的生成代码发现某个改动只影响单边先对照这张表想清楚你的改动属于哪一层。比如你想增加一个循环展开策略CPU 侧可以直接在 CppKernel 循环生成处加GPU 侧则需要考虑 Triton 编译器是否会干扰反而可能画蛇添足。5. 调试两类kernel时踩过的真实坑5.1 如何把生成的cpp和triton代码真正抓出来调试代码生成第一步是把生成代码拿到手。TorchInductor 提供了比较直接的手段在 Python 脚本里设置torch._inductor.config.debug True或者设置环境变量TORCH_LOGSinductor。开启后控制台或临时目录里会出现生成文件的路径文件名通常带torchinductor_前缀。生成的.cpp文件就是 CppKernel 的产物生成的.py文件里可以看到 TritonKernel 的triton.jit函数。这里有个实操技巧不要直接搜所有日志而是把调试开关配合一个极小模型跑。模型越小生成的 kernel 数量越少定位问题越容易。我一般先用torch.compile(lambda x: x 1, dynamicFalse)这种最小的函数打通输出链路确认能抓到生成代码后再逐步替换成真实算子。5.2 坑ACppKernel莫名其妙退化成标量路径有一次我看生成的 CPU 代码发现预期的 AVX2 向量化循环不见了整个 kernel 变成一个简单for (i 0; i n; i) { y[i] x[i] 1; }性能比 eager 模式还差。第一反应是向量化条件没满足。排查过程是这样的先拿TORCH_LOGS确认生成代码然后在控制台里看 Inductor 的 buffer 信息发现输入张量被标记为对齐属性未知。原因是我的输入来自某个自定义的Function返回的扩展张量扩展实现里没有正确设置存储偏移的对齐信息。TorchInductor 的 IR 里虽然有 buffer 结构但alignment属性需要源头提供。解决方式很简单把输入.contiguous()之后再torch.compile对齐信息恢复向量化路径重新出现。这个坑给我的教训是如果你的自定义算子或者第三方扩展涉及内存布局务必先保证张量连续性。否则代码生成器为了安全会主动放弃向量化而不是报错。它选择正确但慢的路径很容易被忽略。5.3 坑BTritonKernel因为shape非2的幂而编译失败另一个印象深刻的坑出现在 GPU 侧。一个对(1000,)张量求sin的小算子编译时直接报错错误信息指向tl.arange不支持非 2 的幂范围。我看生成的 Triton 代码发现tl.arange的范围竟然被写成了tl.arange(0, 1000)而不是tl.arange(0, BLOCK)。原因出在 IR 的 shape 推导某个 pass 把xnumel直接当成了固定值TritonKernel 看到了 Literal type就试图生成一个精确的 arange 而不是用BLOCK做模板参数。tl.arange本身只支持 2 次幂长度所以直接爆炸。最终修复方向是确保TritonKernel生成 arange 时只使用BLOCK这个 constexpr 变量而把真实元素数量放到offsets xnumel的 mask 中比较。这不是 TorchInductor 默认逻辑的问题而是一个自定义 pass 破坏了 constexpr 抽象。它提醒我不要在 IR 阶段随意提前固定 shape尤其不要让 shape 进入生成器的代码模板。5.4 坑Creduction融合后负载不均衡的排查方向还有一个更隐蔽的问题。我在 GPU 上对(1024, 1024)的矩阵做行求和发现生成的 Triton kernel 耗时波动很大。抓出生成代码后看到 reduction 被分成两个 kernel第一个 kernel 做块内 partial reduction第二个 kernel 做最终合并。问题出在第一阶段每个 block 处理的元素数量差异很大导致部分 block 空转。排查这类问题要回到SchedulerNode的 split 策略它没有按 GPU 的调度粒度来划分 block而是沿用了 CPU 侧的循环分块逻辑。GPU 上想解决需要让TritonKernel在生成代码时对 reduction 维度也做类似grid-stride的处理或者重新设计融合边界。这类问题没有一键修复方法但至少能借助生成代码快速定位到是哪一层分配不均。6. 动手改kernel生成器前先想清楚这三件事6.1 明确你要改的是IR、调度还是代码发射很多人看到CppKernel和TritonKernel第一反应是我要改生成代码。但真正要解决的问题往往不在代码发射层。比如你想减少 kernel 数量应该改Scheduler的融合逻辑你想让某个算子走特殊 instruction应该先改 IR 或者注册一个 decompositions你想调向量宽度才是直接改CppKernel的VecSize。如果一上来就改代码发射很容易破坏微妙的不变量。我给一个判断方法先看你的改动是否影响 IR 结构如果影响就尽量在 IR/Scheduler 层做如果只影响最终代码的排版和指令选择才适合在CppKernel/TritonKernel层做。这样可以把风险隔离在某个阶段排查问题也更快。6.2 复用CppKernel/TritonKernel做自研后端的可行路线如果你想基于 TorchInductor 做自己的编译栈不一定要完全重写。比较可行的路线是保留 FX 到 Inductor IR 的降级过程复用Scheduler做融合和 buffer 分配然后替换或继承底层 kernel 生成器。比如继承TritonKernel重写codegen_body把你想要的更细粒度 kernel 结构注入到函数体中。还有一种更轻量的做法定义一个新的torch.compilebackend在接收GraphModule和example_inputs后自己调用torch._inductor.compile_fx或者复用其中部分组件最后返回一个可执行的 wrapper。这样你不需要改动 TorchInductor 内部只是把它当作一个代码生成库来调用。缺点是复用深度有限但胜在稳定。6.3 面对内部API变动时的自我保护TorchInductor 是 PyTorch 2.x 里迭代最快的模块之一内部类名、方法签名经常变动。我自己的习惯是在源码里建一个codegen_ext目录把所有继承/复用的代码集中在里面并且明确声明当前依赖的 PyTorch 版本。每次升级 PyTorch 时先跑一遍codegen_ext下所有继承类的测试确认接口没有破坏。另外要重视日志系统。TorchInductor 的日志非常完善你在二次开发时一定要把自己的代码路径也纳入日志体系否则出了问题很难分辨是上游 pass 还是你的生成器出的错。只要能做到改动隔离 接口测试 日志可追踪基于这两类内核生成器做自定义后端是完全可行的。最后再分享一点个人体会CppKernel和TritonKernel就像是 TorchInductor 的左右手看起来都在生成 kernel但思考方式完全是两套。调 CPU 代码时我会多关注循环结构、对齐和向量化调 GPU 代码时我会多关注 block 大小、mask 和 constexpr 特化。如果你能把上面这些差异内化成直觉再去看cpp.py和triton.py的源码会流畅很多。
