2026/9/11 22:51:02

Kernel 视角:循环、线程与带宽的世界

Kernel 视角:循环、线程与带宽的世界 概述图做全局优化融合、布局、内存复用Kernel 视角看到的是循环、线程、向量化、带宽负责局部优化向量化、tiling、寄存器分配。两者在kernel 的输入输出接口这段话精准地划定了AI编译器中「图编译器Graph Compiler」与「Kernel编译器/算子库Kernel Compiler/Library」之间的职责边界。理解这个边界你就明白了为什么TVM、XLA、MLIR等系统要把编译流程分成两段。编译过程比作「建一栋摩天大楼」图视角 建筑设计师全局规划负责画蓝图、定结构、规划电梯井位置但不负责拧螺丝。Kernel视角 施工队长局部施工负责具体怎么砌墙、怎么扎钢筋但不关心大楼在整个街区的位置。两者的交界点就是「施工蓝图」——也就是代码里的函数签名输入/输出张量的指针、形状、步长、布局图做全局优化融合、布局、内存复用融合把多个算子合并成一个 Kernel省掉中间张量的 DRAM 读写。比如 Conv BN ReLU → 一个 Kernel。布局决定张量在内存里的排列方式比如 NCHW通道连续还是 NHWC空间连续。不同的布局影响后续 Kernel 的内存访问模式。内存复用分析各个中间张量的“存活区间”让不同时间使用的张量共用同一块内存降低峰值显存。小结这些优化都有一个共同点它们只看整张图不看单个算子内部怎么算。所以是“全局”的。Kernel 视角看到的局部优化循环顺序嵌套循环的层次比如是先遍历输出通道还是先遍历空间位置。向量化用 SIMD 指令一次处理多个数据。Tiling分块把大循环切成小块让每块数据适合 L1 缓存 / 寄存器。寄存器分配把频繁使用的中间变量放在寄存器里减少内存读写。两者在“kernel 的输入输出接口”交界图编译器决定 “这个 Kernel 长什么样输入/输出的张量形状、布局、要融合哪些算子”然后把这个任务描述比如“融合 ConvBNReLU输入是 NCHW输出是 NHWCBN 的参数已经折叠好了”交给 Kernel 生成器。Kernel 生成器只负责在这个接口约束下把内部循环写飞。图不关心 Kernel 内部的循环细节Kernel 也不关心图上还有没有其他算子。例子 ConvBNReLU 从图到 KernelInput [N, C, H, W] (NCHW) ↓ Conv2d (stride2, kernel3x3, padding1) ↓ conv_out [N, Co, Ho, Wo] BatchNorm (推理时固定参数) ↓ bn_out [N, Co, Ho, Wo] ReLU ↓ out [N, Co, Ho, Wo]第一步图编译器全局优化图编译器看到整张图做如下决策融合判定Conv 的输出只有一个消费者BNBN 的输出只有一个消费者ReLU。三个算子都是可融合的Conv 是计算密集BN 和 ReLU 是逐元素。判定融合成一个 Kernel。布局选择原本输入是 NCHW。但 Conv 在 NHWC 下内存访问更连续因为卷积是在空间维度滑动通道维度作为内层。图编译器决定将全图布局转为 NHWC并在 Kernel 入口做一次 layout 转换或提前转换输入。内存复用中间张量 conv_out 和 bn_out 原本各占一块内存。融合后它们不需要落 DRAM直接从卷积累加器传给 BN 和 ReLU这些中间值只活在寄存器或 L1 里。图编译器从内存规划中删除这两个 buffer只给最终输出 out 分配内存。BN 参数折叠因为推理时 BN 的 μ、σ、γ、β 固定图编译器计算出折叠后的新权重 w’ 和偏置 b’把 BN 节点从图上删除。最终图编译器输出一个“任务描述”给 Kernel 层Kernel 任务 - 算子Conv BN已折叠 ReLU - 输入张量 A形状 [N, H, W, C] (NHWC)数据类型 fp32 - 权重w 形状 [Co, Kh, Kw, Ci] (NHWC)偏置 b - 输出张量 out形状 [N, Ho, Wo, Co] (NHWC)数据类型 fp32 - 卷积参数stride2, padding1, dilation1 - 要求ReLU 作为 epiloguemax(0, acc b)第二步Kernel 生成器局部优化现在 Kernel 工程师或自动 codegen拿到这个任务开始实现内部细节循环顺序与分块 (Tiling)输出是 [N, Ho, Wo, Co]需要 5 层循环N, Ho, Wo, Co, 累加 Ci 和 Kh/Kw。为了利用 L1 缓存他们决定对输出空间维度 (Ho, Wo) 和输入通道 Ci 做分块。比如每次处理 8x8 的输出像素同时只加载需要的 Ci 分块比如 32 个通道。向量化 (SIMD)内层累加循环中对 Ci 的累加可以使用 SIMD 指令如 AVX2 的 vfmadd231ps一次处理 8 个浮点数。同时ReLU 的 max(0, acc) 也可以用向量比较指令一条完成。寄存器分配累加器 acc 是寄存器中的临时变量如 ymm0。分块后每个线程/核心负责一小块输出所需的权重和输入数据尽量放在 L1/L2 缓存避免反复从 DRAM 读取。并行线程如果是 GPU每个线程块处理 (Ho, Wo) 的一个大块每个线程处理一个输出元素。如果是 CPU用 OpenMP 并行化外层 N 或 Co 循环。最终 Kernel 的伪代码简化// 融合 Kernel: fused_conv_bn_relu_nhwcfor(intn0;nN;n){for(intho0;hoHo;hoTILE_H){for(intwo0;woWo;woTILE_W){for(intco0;coCo;coTILE_C){// 加载权重分块到 L1// 对每个输出像素 (hodh, wodw) 累加for(intci0;ciCi;ci){// 向量化累加acc[0..7]input[...]*weight[...];}// 加上折叠后的偏置应用 ReLUout[n,hodh,wodw,codc]max(accb[codc],0);}}}}接口在哪里图编译器传下来的东西输入张量地址、权重地址、输出地址、形状参数 (N, H, W, C, Ho, Wo, Co, stride, padding)、折叠后的偏置。Kernel 内部完全不知道图里还有没有别的算子也不关心 BN 参数是怎么折叠的——它只按照接口把循环跑好。总结图编译器负责“宏观规划”决定哪些算子打包成一个 Kernel数据怎么排放中间结果怎么省掉。Kernel 生成器负责“微观落地”在给定的任务描述下用最优的循环、分块、向量化、寄存器分配把它实现出来。两者通过 Kernel 的函数签名输入输出张量、参数 这个明确接口解耦这也是现代 AI 编译器如 TVM、XLA、MLIR能够同时兼顾全局优化和局部极致性能的根本原因。