第5章：XLA 与 JIT 编译——从图到机器码

---

5.1 一个让你好奇的数据

在第 2 章，你跑了这个测试：

f(x) = x⁴ + 1, 100000 次调用
SimpleGo 后端：约 759ms（132K runs/sec）

如果你在 Linux 上用 XLA 后端跑同样的代码，结果大约是 SimpleGo 的 3-5 倍。同一段图定义，同一个数学公式，速度为什么差这么多？

答案在于编译器的质量。SimpleGo 把图翻译成 Go 函数调用链（加上 SIMD 加速），而 XLA 做了一整套更激进的编译优化。这一章讲"那一整套"里到底包含什么。

---

5.2 JIT 编译：不是预编译，不是解释执行

先理清一个基础概念。

Go 程序是预编译（AOT，Ahead-of-Time）的：你敲 go build，编译器把 Go 代码提前翻译成机器码。程序启动时，所有代码已经是机器码了。

Python 程序是解释执行的：代码一行一行被解释器读取、翻译、执行。

GoMLX 用的是第三种方式——JIT（Just-in-Time，即时编译）：你的图定义代码在 Go 里（AOT 编译），但图本身的执行代码，是在运行到 MustExec 的那一刻才编译成机器码的。

exec := graph.MustNewExec(backend, graphFn) // ← 图构建（Go 代码，AOT 编译的）
result := exec.MustExec(3.0)                 // ← 首次调用：JIT 编译 + 执行
result2 := exec.MustExec(5.0)                // ← 后续调用：直接执行（命中缓存）

为什么用 JIT？因为它有一项 AOT 没有的优势：它知道运行时信息。它知道 x 的具体形状是 [128, 512] 而不是 [?, 512]。知道了精确形状，就可以做形状特化优化。

---

5.3 XLA 的完整编译流水线

XLA（Accelerated Linear Algebra）是 Google 开发的开源线性代数编译器。它自己不运行数学——它编译数学计算。

GoMLX 的 XLA 后端走这条流水线：

GoMLX 计算图（DAG）
       │
       ▼
  ┌─────────────┐
  │ HLO（高级IR） │  ← 平台无关的线性代数中间表示
  └─────────────┘
       │
       ▼
  ┌─────────────┐
  │   XLA 优化器  │  ← 算子融合、布局选择、代数化简……
  └─────────────┘
       │
       ▼
  ┌─────────────┐
  │   LLVM IR    │  ← 通用的低级中间表示
  └─────────────┘
       │
       ▼
  ┌─────────────┐
  │   机器码      │  ← x86-64 / CUDA PTX / ARM NEON
  └─────────────┘

一层一层走。

第一层：GoMLX 图 → HLO。GoMLX 把你的 graph.Add(a, b) 翻译成类似于 %add = add(%a, %b) 的 HLO 指令。翻译是保语义但失细节的。

第二层：HLO → 优化后的 HLO。XLA 在 HLO 层面做硬件无关的优化：代数化简、常量折叠、DCE（死代码消除）。

---

5.4 算子融合：XLA 的杀手级优化

看这个例子：

func fusedOps(x *graph.Node) *graph.Node {
    a := graph.MulScalar(x, 2.0)   // 乘法
    b := graph.AddScalar(a, 1.0)   // 加法
    c := graph.Relu(b)             // ReLU 激活
    return c
}

不融合时，三次 GPU kernel 启动 + 三次显存分配/释放 + 两次中间结果的显存读写。

算子融合后，XLA 把这三个运算合并成一个 kernel：一次 kernel 启动、零中间显存分配、零中间数据搬运。GPU 上 kernel 启动和显存带宽才是真开销——计算反而次要。

XLA 的融合规则：逐元素运算之间的融合最激进；缩减运算可以和后续逐元素融合；矩阵乘法的融合更保守。

对读者的工程意义：当你写 GoMLX 代码时，连续的逐元素运算不会产生中间张量——XLA 帮你在编译期融合掉了。你可以放心写清晰、分步的公式。

---

5.5 LLVM IR → 机器码

XLA 优化完 HLO 后，把它扔给 LLVM。LLVM 是工业标准的编译器基础设施。

LLVM 拿到 XLA 优化后的 HLO，先翻译成 LLVM IR（一种低级的、接近汇编的中间表示），然后做自己的优化：指令调度、寄存器分配、循环展开……最后生成目标平台的机器码。

对于 CPU：生成 x86-64 或 ARM NEON 指令。

对于 GPU：生成 CUDA PTX（NVIDIA）或 SPIR-V，然后再由 GPU 驱动编译到 GPU 指令。

经过 XLA + LLVM 两重编译优化后，最终生成的代码质量，理论上可以逼近手写 CUDA kernel。

---

5.6 SimpleGo 后端 vs XLA 后端：两种编译哲学

SimpleGo 没有 XLA。它走另一条路：

GoMLX 计算图
    │
    ▼
Go 函数调用链（手工优化的，含 SIMD）
    │
    ▼
Go 编译器（AOT，已包含在 go build 中）
    │
    ▼
机器码

SimpleGo 的"编译"本质上是用 Go 代码模拟了每一个算子的执行。它的优势不是快，而是： 1. 零外部依赖：不需要 XLA、CUDA、PJRT 2. 全平台：Windows、macOS、Linux，甚至 WASM 3. SIMD 加速：核心算子用 AVX-2/AVX-512 手写汇编

对于中小模型，SimpleGo 每秒 13 万次调用，处理大多数 CPU 推理场景绰绰有余。但它没有算子融合——这也解释了 SimpleGo 为什么比 XLA 慢 3-5 倍。

---

5.7 编译缓存：为什么第一次慢，之后快

第一次调用触发了完整的编译流水线（图构建 + XLA 优化 + LLVM 编译）。第二次命中缓存——直接跑编译好的机器码。

两个数字值得记住：编译一个中等复杂的图（如 ResNet-18 的一个 forward pass）在 XLA 上大约需要 0.1-2 秒。相比训练几小时，这几乎可以忽略。但如果你每秒钟创建上百个新形状的图，编译开销就会成为瓶颈。

---

5.8 收束：编译的全景

第 5 章的核心认知：

> JIT 编译 = 运行时编译，获得精确形状信息 → XLA 做激进优化（尤其是算子融合）→ 生成高性能机器码。编译一次，执行无数次。缓存命中后，执行开销接近零。

---

> 作者：步子哥，资深AI研究员 | 连载中：《GoMLX从入门到精通》 > 上一章：第4章计算图核心数据结构 > 下一章预告：第6章张量——数据的形状与流动