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 长什么样（输入/输出的张量形状、布局、要融合哪些算子）”

然后把这个任务描述（比如“融合 Conv+BN+ReLU，输入是 NCHW，输出是 NHWC，BN 的参数已经折叠好了”）交给 Kernel 生成器。Kernel 生成器只负责在这个接口约束下，把内部循环写飞。

图不关心 Kernel 内部的循环细节，Kernel 也不关心图上还有没有其他算子。

例子： Conv+BN+ReLU 从图到 Kernel

Input [N, C, H, W] (NCHW) ↓ Conv2d (stride=2, kernel=3x3, padding=1) ↓ conv_out [N, Co, Ho, Wo] BatchNorm (推理时固定参数) ↓ bn_out [N, Co, Ho, Wo] ReLU ↓ out [N, Co, Ho, Wo]

第一步：图编译器（全局优化）

图编译器看到整张图，做如下决策：

融合判定

Conv 的输出只有一个消费者（BN），BN 的输出只有一个消费者（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 - 卷积参数：stride=2, padding=1, dilation=1 - 要求：ReLU 作为 epilogue（max(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(intn=0;n<N;++n){for(intho=0;ho<Ho;ho+=TILE_H){for(intwo=0;wo<Wo;wo+=TILE_W){for(intco=0;co<Co;co+=TILE_C){// 加载权重分块到 L1// 对每个输出像素 (ho+dh, wo+dw) 累加for(intci=0;ci<Ci;++ci){// 向量化累加acc[0..7]+=input[...]*weight[...];}// 加上折叠后的偏置，应用 ReLUout[n,ho+dh,wo+dw,co+dc]=max(acc+b'[co+dc],0);}}}}