ML Compilers Overview

01ML Compiler Overview#

ML Compiler 位于深度学习框架和底层硬件之间，作用是把 TensorFlow、PyTorch、JAX 等框架描述的高层张量计算转换成能在 CPU/GPU/TPU 等硬件上高效执行的程序。

它的典型架构是 Frontend → IR → Optimization → Backend:

• Frontend 接收框架已经构造或捕获出来的计算图/算子程序（例如 tf.function 得到的 TF Graph、PyTorch Dynamo 得到的 FX Graph），转换成统一的中间表示 IR
• 中间层进行算子融合、常量折叠、布局变换、内存规划、并行化等与机器学习计算相关的优化
• Backend 再根据目标硬件进行 lowering、调度和代码生成

因而 ML compiler 输入通常是带有 tensor shape、dtype、算子及数据依赖关系的计算图或 IR，输出则是面向特定设备的低层 IR、kernel 或可执行程序。

02IR Pipeline#

下面这套 IR 划分描述的是一个典型 ML Compiler 从高层模型计算逐渐 lowering 到硬件可执行 kernel 的过程。不同编译器实际使用的 IR 名称和数据结构可能不同，但抽象层次通常可以归纳为：

下面各层的 Example 来自同一次 Emmy compile --ir 运行。输入程序为 SiluAndMul(x)：沿最后一维将 float32[1, 32, 4096] 的 x 分成两半，对前半部分计算 SiLU，再与后半部分逐元素相乘。导出环境为 PyTorch 2.13.0、Emmy c6755ee，目标 GPU 为 A6000

class SiluAndMul(nn.Module):
    def forward(self, x):
        left, right = x.chunk(2, dim=-1)
        return F.silu(left) * right

2.1Framework IR#

Framework IR 表达模型中的高层算子以及算子之间的数据依赖关系

这一层关注的是“要计算什么”，通常仍然保留深度学习框架中的高层 tensor operator，例如 matrix multiplication、convolution、attention、softmax 等，而不会描述这些算子内部如何实现

Example

Emmy 的 Torch IR 仍接近捕获到的 PyTorch operator：

inputs:
  x: (1, 32, 4096) f32

chunk_0 = slice(x, dim=2, start=0)     -> (1, 32, 2048) f32
chunk_1 = slice(x, dim=2, start=2048)  -> (1, 32, 2048) f32
silu    = silu(chunk_0)                -> (1, 32, 2048) f32
mul     = multiply(silu, chunk_1)      -> (1, 32, 2048) f32

outputs:
  mul: (1, 32, 2048) f32

chunk(2, -1) 已被静态解析为两个 slice，但 silu 和 multiply 仍是高层 tensor operator。这一层描述两个 slice、SiLU 和 multiply 的数据依赖，尚未决定 slice 是否复制数据，也没有显式 loop 或 hardware mapping。

2.2Tensor IR#

Tensor IR 将高层算子分解成更基础、统一的张量计算语义。 对大量规则的 dense tensor computation，可以把其核心抽象归纳为三类：

1. Elementwise：描述每个输出位置执行的标量计算，例如 add、mul、exp
2. Reduction：描述沿某个维度进行聚合，例如 sum、max
3. IndexMap：使用坐标映射统一表示 broadcast、reshape、transpose、slice 等布局和索引变换

Example

同一个 SiluAndMul 在 Emmy Tensor IR 中被表示为 IndexMap 与 Elementwise primitive：

chunk_0      = x[i, j, k]
chunk_1      = x[i, j, k + 2048]
silu_neg     = negative(chunk_0)
silu_exp     = exp(silu_neg)
silu_denom   = add(1.0, silu_exp)
silu_sigmoid = reciprocal(silu_denom)
silu         = multiply(chunk_0, silu_sigmoid)
mul          = multiply(silu, chunk_1)

• 两个 slice 已经变成直接作用于原输入的坐标映射，因此 graph 中没有 chunk copy
• 高层 silu 则被 decomposition 为 negative、exp、add、reciprocal 和 multiply
• 不同框架中的高层 operator 正是通过这种 canonicalization 被统一到少数基础 tensor semantics，供后续 Loop IR 使用。

下图把这两件事画在一起：x 的两半没有被复制，而是由两个 IndexMap 直接按坐标读取；虚线框内是 silu decomposition 后的 elementwise 链。

2.3Loop IR#

Loop IR 将 Tensor IR 中的张量计算显式展开成 iteration、indexing、load/store 和 reduction。

Tensor IR 中的 primitive 会先被 lifting 为 loop nest，再在合法的情况下进行 loop fusion，将相邻计算合并为尽可能少的 kernel loop nests。

Example

Emmy 将前面的 IndexMap 和 Elementwise nodes lifting 并 fusion 为一个 pointwise loop nest：

for a0 in 0..32:
    for a1 in 0..2048:
        left = load x[0, a0, a1]
        neg = negative(left)
        exp_neg = exp(neg)
        sigmoid = reciprocal(add(1.0, exp_neg))
        silu = multiply(left, sigmoid)
        right = load x[0, a0, a1 + 2048]
        output = multiply(right, silu)
        store mul[0, a0, a1] = output

a0 和 a1 都是 free / parallel axes。两个 chunk load、SiLU decomposition 和最终 multiply 已进入同一次 loop iteration，没有 intermediate tensor store。

以某一个 (a0, a1) 为例，一次 iteration 从 x 的同一行读取前后两半中对应的两个元素，在 loop body 中完成全部计算，再写出一个输出元素：

此时 IR 已明确 iteration space、indexing 和 load/store，但还没有决定展开因子、worker 数量或 CUDA thread mapping。

2.4Tile IR#

Tile IR 在 Loop IR 的基础上加入 schedule 信息，描述循环如何被切分、重排、映射到并行计算资源，以及数据如何在内存层级之间 staging。

在 ML Compiler 中，schedule 指的是：

典型的 scheduling 决策包括：

1. Tiling：把大的 iteration space 划分成较小的 tile
2. Reordering：调整 loop 的执行顺序
3. Parallelization：将计算轴抽象地绑定到 THREAD、BLOCK 等并行维度
4. Memory staging：决定哪些数据需要经过 shared memory 等片上存储
5. Hardware-aware scheduling：根据目标硬件加入规约、同步和流水线等调度决策

Example

本次默认 pointwise schedule 选择 4 路展开。下面保留第一个和最后一个 unrolled lane：

in1__u0 = load x[0, a0, a1 * 4 + 0]
in2__u0 = load x[0, a0, a1 * 4 + 0 + 2048]
v0__u0 = negative(in1__u0)
v1__u0 = exp(v0__u0)
v2__u0 = add(1.0, v1__u0)
v3__u0 = reciprocal(v2__u0)
v5__u0 = multiply(in2__u0, multiply(in1__u0, v3__u0))

# __u1 和 __u2 执行相同计算

in1__u3 = load x[0, a0, a1 * 4 + 3]
in2__u3 = load x[0, a0, a1 * 4 + 3 + 2048]
v0__u3 = negative(in1__u3)
v1__u3 = exp(v0__u3)
v2__u3 = add(1.0, v1__u3)
v3__u3 = reciprocal(v2__u3)
v5__u3 = multiply(in2__u3, multiply(in1__u3, v3__u3))

store mul[0, a0, a1 * 4 + 0] = v5__u0
store mul[0, a0, a1 * 4 + 1] = v5__u1
store mul[0, a0, a1 * 4 + 2] = v5__u2
store mul[0, a0, a1 * 4 + 3] = v5__u3

一个逻辑 worker 现在处理 a1 * 4 + {0, 1, 2, 3} 四个连续位置。__u0 到 __u3 是四个彼此独立的 unrolled lane；前后两个 chunk 的 load 使用相同展开因子，为后续 vector store 准备连续结果。

2.5Kernel IR#

Kernel IR 将 Tile IR 中的抽象调度决策物化为具体的硬件 primitive。 这一层已经面向目标硬件的执行模型，但仍然是一种结构化 IR，而不是最终的 CUDA、Triton 或 C++ 源代码。

经过 Tile IR lowering 之后，Kernel IR 通常已经明确：

1. kernel 的输入、输出和边界
2. block、thread 等具体 launch geometry
3. global、shared、local 等内存空间
4. indexing、load、store 和异步拷贝方式
5. barrier、同步和跨线程规约操作
6. 目标相关的并行工作分解

Example

Emmy Kernel IR 将 4 路展开落实为 worker tile、具体 f32 scalar 和 vector store。下面只展开第一个 lane：

Tile[a0, a1] (N=16384)
    in1__u0 = load x[0, a0, a1 * 4 + 0]
    f32 v0__u0 = negative(in1__u0)
    f32 v1__u0 = exp(v0__u0)
    f32 v2__u0 = add(1.0, v1__u0)
    f32 v3__u0 = reciprocal(v2__u0)
    f32 v4__u0 = multiply(in1__u0, v3__u0)
    in2__u0 = load x[0, a0, a1 * 4 + 0 + 2048]
    f32 v5__u0 = multiply(in2__u0, v4__u0)

    # __u1、__u2、__u3 采用相同结构
    mul[0, a0, a1 * 4] = (v5__u0, v5__u1, v5__u2, v5__u3)

N=16384 对应 32 × 512 个逻辑 worker，每个 worker 处理四个连续元素，因此覆盖 32 × 2048 个输出。四个结果被组合为一次 vector store。该 kernel 没有跨 worker dependency，所以不需要 reduction、shared memory、barrier 或 synchronization。

Kernel IR 可以理解为已经形成 kernel boundary，并将 schedule 落实为 target-oriented worker、dtype、memory operation 和 synchronization primitive，但尚未渲染成目标语言源码。

不同 compiler 的分层边界并不完全相同：Emmy 在这一层使用 Tile[a0, a1] 表达 worker grid，具体 blockIdx.x、threadIdx.x 和 block size 在 CUDA rendering 时才完全物化。

2.6Target Code#

Target Code 将 Kernel IR 渲染为目标 backend 可以继续编译的源代码。 它已经使用目标语言的 thread index、pointer arithmetic、intrinsic 和 vector type，但通常还不是 GPU 直接执行的 machine code。

Example

Emmy 将上述 Kernel IR 渲染为 CUDA C++。下面保留 thread mapping、四个 unrolled lanes 和最终 vector store：

extern "C" __global__
__launch_bounds__(256)
void k_slice_pointwise_dcfb6d(
    const float * x,
    const float * silu_one,
    float * mul
) {
    int gid = blockIdx.x * blockDim.x + threadIdx.x;
    int a0 = gid / 512;
    int a1 = gid % 512;

    float left_u0 = x[a0 * 4096 + a1 * 4];
    float right_u0 = x[a0 * 4096 + a1 * 4 + 2048];
    float out_u0 = right_u0 * (
        left_u0 * (1.0f / (silu_one[0] + expf(0.0f - left_u0)))
    );

    float left_u1 = x[a0 * 4096 + a1 * 4 + 1];
    float right_u1 = x[a0 * 4096 + a1 * 4 + 1 + 2048];
    float out_u1 = right_u1 * (
        left_u1 * (1.0f / (silu_one[0] + expf(0.0f - left_u1)))
    );

    float left_u2 = x[a0 * 4096 + a1 * 4 + 2];
    float right_u2 = x[a0 * 4096 + a1 * 4 + 2 + 2048];
    float out_u2 = right_u2 * (
        left_u2 * (1.0f / (silu_one[0] + expf(0.0f - left_u2)))
    );

    float left_u3 = x[a0 * 4096 + a1 * 4 + 3];
    float right_u3 = x[a0 * 4096 + a1 * 4 + 3 + 2048];
    float out_u3 = right_u3 * (
        left_u3 * (1.0f / (silu_one[0] + expf(0.0f - left_u3)))
    );

    float4 output = make_float4(out_u0, out_u1, out_u2, out_u3);
    *reinterpret_cast<float4 *>(&mul[a0 * 2048 + a1 * 4]) = output;
}

Tile[a0, a1] 被映射为一维 gid，再通过 / 512 和 % 512 恢复两个 worker axes。16384 个 worker 按 256 threads/block 启动时需要 64 个 block。输入仍是 scalar load，明确的 vectorization 出现在最终 float4 store。此次导出停在 CUDA source；还需经过 nvcc 等工具链才能得到 PTX、cubin 或其他 machine-code artifact。

以 blockIdx.x = 5、threadIdx.x = 37 这个线程为例，它从 gid 还原出 (a0, a1)，最终写出输出第 2 行的 4 个连续元素：

2.7Summary#

整个 lowering 过程可以概括为：

因此各层最核心的区别可以理解为：

从 ML Compiler 的角度看，这整个过程本质上就是不断进行 lowering：

也就是不断从“描述计算是什么”下降到“描述硬件具体怎么执行这个计算”。

2.8Summary#

Compiler 的典型架构是 Frontend → IR → Optimization → Backend Code:

• Frontend 接收框架已经构造或捕获出来的计算图/算子程序（例如 tf.function 得到的 TF Graph、PyTorch Dynamo 得到的 FX Graph），转换成统一的中间表示 IR
• 中间层进行算子融合、常量折叠、布局变换、内存规划、并行化等与机器学习计算相关的优化
• Backend 再根据目标硬件进行 lowering、调度和代码生成

03PyTorch Eager Mode#

Eager mode 下不存在 “Python → IR → Backend Code” 这个编译过程

Python 每执行一个 PyTorch op，就通过 Dispatcher 找到已经实现好的 CUDA kernel，然后直接 launch

整体流程如下:

其中第 3 步的 Dispatcher 并不是一次查表。每个 Tensor 带有一组 dispatch key，Dispatcher 从优先级最高的 key 开始处理，每一层做完自己的工作后 redispatch 到下一层。以 requires_grad=True 的 CUDA tensor 执行 torch.mul(a, b) 为例：

在 Eager Mode 下，CUDA code 大部分是在 PyTorch / CUDA libraries 构建时就编译好了

这些 kernel 彼此独立，所以 eager 无法跨 op 优化。以前文的 SiluAndMul 为例（每个 tensor 有 N 个元素），eager 需要两个 kernel，silu 的结果必须先写回 HBM 再被 mul 读出；torch.compile 生成的 fused kernel 把中间值留在 register 中，HBM 传输从 5N 降到 3N，launch 也从 2 次降到 1 次：

04PyTorch torch.compile#

4.1Overview#

总的来说可以分为四层

4.2Workflow#

1 Graph Capture：TorchDynamo#

Dynamo 会 hook CPython 的 frame evaluation，解释 Python bytecode，把其中能编译的 Tensor 操作提取成 torch.fx.GraphModule

相关的操作有:

• Python bytecode tracing
• 捕获 PyTorch Tensor 操作
• 建立 FX Graph
• 处理 Python control flow
• Shape / dtype / stride 等 symbolic 信息
• 建立 guards
• 遇到不能 capture 的代码时产生 graph break

Dynamo 通过 PEP 523 接管 frame evaluation。每个 code object 带有一个 cache，每个 cache entry 都有一组 guard；调用时先检查 guard，命中就直接执行编译结果，未命中才重新 trace 和编译：

遇到无法 trace 的代码（例如 print）时，Dynamo 会把已经捕获的部分编译成一张图，让 CPython 执行这段代码，再为剩余部分生成 __resume_at_N 函数。这个函数会在新的 frame 中再次被 Dynamo 捕获：

2 Graph Normalize: AOTAutograd#

这一阶段主要负责对 Dynamo 捕获的高层 PyTorch graph 进行 functionalization、decomposition 以及 alias/mutation 规范化，生成适合后续后端优化与 lowering 的 ATen FX Graph

典型工作包括：

• Functionalization

  • 尽量处理 mutation / inplace / alias

• Decomposition

  • 复杂 op 拆成 backend 更容易处理的 ATen op
• 规范化 graph
• 根据 inference/training 情况 dispatch 给后端 compiler

3 Graph Optimization：TorchInductor Graph Passes#

Inductor 会先做一系列 FX-level optimization/pass，比如:

4 Lowering + Scheduling: TorchInductor IR#

Inductor 会从 FX / ATen graph lower 到自己的 loop-level IR

随后 scheduler 会决定：

• 哪些 op 可以 fuse
• kernel 边界
• execution order
• buffer dependencies
• memory reuse / planning
• tile / scheduling strategy

5 Backend Codegen#

PyTorch 2.13：CuTeDSL Native DSL Backend#

PyTorch 2.13 的一个重点是继续扩展 TorchInductor 的 GPU backend。对于 NVIDIA GPU，如果暂时忽略 external kernel，过去常将 Inductor 的主要 codegen path 简化为：

从 PyTorch 2.13 开始，Inductor 增加了 CuTeDSL 这条 alternative code-generation path：

CuTeDSL 是基于 NVIDIA CuTe（CUDA Templates）的 Python-native DSL，可以直接表达 GPU tensor layout、tiling strategy 和 memory access pattern。当前 Inductor 集成主要面向 GEMM 与 RMSNorm 等关键 operation，为这些 workload 提供 Triton 之外的另一条高性能 codegen path；它不是对所有 Inductor GPU kernel 的整体替换。

PyTorch 2.13 还将相关 kernel compilation 从 thread pool 移到 subprocess pool，以绕开 Python GIL 对并行编译的限制。官方将这项能力描述为兼顾高性能与更快 compilation 的第二条路径，但目前仍标记为 API Unstable，实际是否选择 CuTeDSL 取决于 operation、hardware、dtype、layout 和 backend support。

在 NVIDIA GPU 上，默认仍是 Triton codegen path，对应 pipeline 为：