← 返回主题列表
Q
QianXun
@QianXun · 2026年06月13日 09:20 · 0浏览

《GoMLX从入门到精通》第3章：GoMLX的世界观——计算图、后端与执行模型

第3章：GoMLX 的世界观——计算图、后端与执行模型

---

3.1 一个让你困惑的瞬间

在第 2 章，你写了这段代码：

addGraph := func(a, b *graph.Node) *graph.Node {
    return graph.Add(a, b)
}
exec := graph.MustNewExec(backend, addGraph)
results := exec.MustExec(1.0, 2.0)
fmt.Println(results[0].Value()) // 3

一切正常。但你不妨试试这个：

counter := 0
weirdGraph := func(x *graph.Node) *graph.Node {
    counter++ // 我在图定义函数里修改了外部变量！
    return graph.AddScalar(x, float64(counter))
}

exec := graph.MustNewExec(backend, weirdGraph)
r1 := exec.MustExec1(1.0) // counter = 1，返回 1 + 1 = 2 ？
r2 := exec.MustExec1(1.0) // counter = 2，返回 1 + 2 = 3 ？

fmt.Println(r1.Value()) // 实际输出：2
fmt.Println(r2.Value()) // 实际输出：2？？？为什么还是2？

你可能会愣住——counter++ 明明是写在图定义函数里的，而且 MustNewExec 确实调用了它一次（counter 从 0 变成 1），为什么第二次 MustExec1 时 counter 没再增加？

答案，就是这整章要讲的东西。

---

3.2 命令式编程 vs 符号式编程

Go 程序员（以及几乎所有的命令式语言程序员）的习惯是：写一行代码，它就做一件事。

a := 1
b := 2
c := a + b  // 这行代码执行时，a+b 立刻被计算，c 立刻被赋值
a = 5       // 修改 a
// c 还是 3——因为 c = a + b 执行时 a 还是 1

这叫命令式编程（Imperative Programming）。代码的每一行都是一个即时执行的指令。

GoMLX 不是这样工作的。你写的图定义函数不是执行计算，而是描述计算。它属于符号式编程（Symbolic Programming）——你先"画"一张完整的计算图，然后编译器一次性地编译并执行整张图。

3.1 中那个诡异现象的答案就在这里：

counter := 0
weirdGraph := func(x *graph.Node) *graph.Node {
    counter++ // 这行只在图构建阶段执行！只执行一次！
    return graph.AddScalar(x, float64(counter))
}

exec := graph.MustNewExec(backend, weirdGraph)
// ↑ 此时调用 weirdGraph 一次：counter 变成 1，图上记录了 "x + 1.0"
//   编译完成，图被缓存。

r1 := exec.MustExec1(1.0) // 执行编译好的图：1.0 + 1.0 = 2.0
// counter 没有变——这行只是跑机器码，不跑 Go 代码

r2 := exec.MustExec1(1.0) // 同一张图，同样的机器码：1.0 + 1.0 = 2.0
// counter 还是 1——图早就编译好了，不会再调 weirdGraph

一张表说清区别：

命令式编程（Go 默认）符号式编程（GoMLX）
你的代码描述的是做什么、怎么做、按什么顺序做数据之间的计算关系
代码执行时机写一行，执行一行分批：先构建图（一次），再编译+执行（多次）
变量的含义内存中的值，随时可变图中的节点，编译后固化
"加法"的语义a + b → 立即读 a 和 b 的值，算和graph.Add(a, b) → 在图里画一个加法节点，表明 a 和 b 是这个节点的输入
做菜的类比：命令式编程是你亲自站在锅前，切菜、下锅、翻炒、加盐——每一步都即时发生。符号式编程是你写一份食谱，交给一个专业的厨师团队（JIT 编译器），他们读完整份食谱后才开始高效地做菜——你以为自己在"炒菜"，其实你一直在"写食谱"。

---

3.3 两阶段生命周期：Build → Execute

理解 GoMLX，只需要把它拆成两个阶段来看。

阶段一：图构建（Graph Building）

这是你的 Go 代码真正被执行的阶段。发生在 MustNewExec 第一次调用你的图定义函数时。

在这个阶段：

  • 图定义函数被调用了一次（或更准确地说，每种输入形状的组合各调用一次——见 3.5 节）
  • 函数内部的 Go 代码（counter++、if 分支、for 循环）全部正常执行
  • 但是——任何 graph.Add、graph.Square、graph.MulScalar 等 GoMLX 运算 API，做的事情不是"算数值"，而是往图上添加节点
  • 最终产物是一张 DAG（有向无环图），每个节点代表一种运算
你可以把它想象成"录音阶段"。你在麦克风前面说话，声音被录进去了——但没人能听到。你只是在生成一份"待播放清单"。

阶段二：图执行（Graph Execution）

这是你的图被机器真正执行的阶段。发生在每次 exec.MustExec(args...) 调用时。

在这个阶段：

  • 你的 Go 代码不再运行。 运行的只是编译好的机器码（SimpleGo 后端的 Go 原生代码，或 XLA 后端的 LLVM JIT 产物）
  • 图里的每个节点，从输入参数开始，按依赖顺序依次被计算
  • 有 GPU 时，整张图在显存里跑——数据不离开 GPU
  • 结果从设备内存搬回 Go 一侧
这就像"播放阶段"。你按下播放键，录好的声音被放大、混音、从扬声器放出来。这个过程中，麦克风是静音的——你不再说话，只是听。

核心推论

一旦你接受了"构建→执行"的两阶段模型，很多表层现象立刻就有了解释：

为什么修改外部变量不会影响执行结果？ 因为执行阶段只跑机器码，不跑你的 Go 代码。

为什么 MustNewExec 很快但 MustExec 可能慢？ 因为图构建可能做一些复杂逻辑（Go 代码层面），而图执行只是纯粹的后端运算（编译后的高效代码）。反过来也可能：简单图的构建很快，但首次 MustExec 要触发 JIT 编译（慢），后续调用直接跑机器码（快）。

为什么 GoMLX 能换后端？ 因为图定义（你的 Go 代码）不包含任何后端特定逻辑。它只产出抽象的运算节点。后端拿到这张图后，各自决定怎样把它翻译成自己的可执行格式。XLA 后端翻译成 HLO→LLVM IR→机器码；SimpleGo 后端翻译成 Go 函数调用链。

---

3.4 Node、Exec、Backend

GoMLX 的世界里，有三个角色一直在互动。理解它们的职责边界，你就理解了 GoMLX 的架构。

*graph.Node——图中的节点

// 每次调用 graph.Add, graph.Square 等，都生成一个新的 *Node
xSquared := graph.Square(x)     // *Node：代表 x² 这个运算
twoX := graph.MulScalar(x, 2.0) // *Node：代表 2x 这个运算
sum := graph.Add(xSquared, twoX) // *Node：代表 x² + 2x 这个运算

*graph.Node 是什么？

  • 它是计算图中的一个节点，代表一个运算的结果
  • 它携带的信息包括：做了什么运算（Add/Square/MulScalar...）、输入是哪些节点、输出的形状和数据类型
  • 它不包含任何数值——直到图执行阶段
  • 一个 Node 可以被多个下游 Node 引用（这就是"共享子表达式"的来源——如果同一个 Node 被引用了三次，底层只会算一次）

*graph.Exec——编译后的可执行图

exec := graph.MustNewExec(backend, graphFn)

*graph.Exec 是什么？

  • 它是图的编译结果 + 执行接口
  • 内部维护了一个"形状→编译图"的缓存（默认最多 32 个不同形状版本）
  • 首次调用 MustExec(args...) 时，查看 args 的形状，如果没有对应的编译版本，就触发一次全新的图构建 + 编译——这就是"慢"的来源
  • 后续相同形状的调用直接命中缓存——这就是"快"的来源
  • 线程安全——可以并发调用 MustExec

backends.Backend——计算后端

backend, _ := backends.New()

backends.Backend 是什么？

  • 它是计算设备的抽象。CPU、GPU、TPU，都是 Backend
  • 在图定义阶段完全不参与——你不需要知道它在哪
  • 在图编译阶段介入——它负责把抽象的图翻译成自己的可执行格式
  • 在图执行阶段全权负责——分配设备内存、调度运算、搬运数据
  • 不同的 Backend，编译出的机器码完全不同——但你的图定义代码完全一样
三个角色的协作关系，一句人话总结：

> ***graph.Node 是你画的图纸，*graph.Exec 是施工队的翻译和执行总控，backends.Backend 是具体的施工设备（手工具、吊车、挖掘机）。

---

3.5 静态形状限制：设计选择，不是缺陷

你有没有想过：为什么同一个 exec 不能处理不同形状的输入？

exec := graph.MustNewExec(backend, func(x *graph.Node) *graph.Node {
    return graph.Square(x)
})

// ✓ 工作正常
r1 := exec.MustExec1([]float64{1.0, 2.0, 3.0}) // shape: [3]

// ✗ 形状变了，需要重建图
r2 := exec.MustExec1([]float64{1.0, 2.0}) // shape: [2] → 新的图构建+编译

这不是 GoMLX 的 bug。这是刻意为之的性能优化策略**。

编译期形状特化（Shape Specialization）

编译器在知道所有张量的确切形状时，可以做一系列"如果不知道就不敢做"的优化：

  • 循环展开：一个长度为 128 的向量加法，可以直接展开成 128 条标量加法的流水线
  • 内存布局优化：提前计算每个中间张量的内存偏移，消除运行时地址计算
  • 算子融合：把 Mul → Add → Relu 融合成一个 Kernel（第 5 章会展开）
这些优化都需要在编译时就知道形状。如果运行时形状可变，编译器就只能生成"通用"代码——通用意味着慢。

再加上缓存策略：*graph.Exec 内部缓存了最多 32 个不同形状的编译版本。如果你在一个循环里交替调用 shape=[128] 和 shape=[64] 的输入，第一次各慢一次（编译），之后每次都快（命中缓存）。

这个权衡合理：

exec := graph.MustNewExec(backend, func(x *graph.Node) *graph.Node {
    return graph.Square(x)
})

// 批次大小在 [64, 128, 256] 之间轮流出现——每种第一次慢，之后快
for _, batch := range batches {
    result := exec.MustExec1(batch) // 第 1 批 64：编译（慢）
                                    // 第 1 批 128：编译（慢）
                                    // 第 2 批 64：命中缓存（快）
                                    // 第 2 批 128：命中缓存（快）
}

实践中怎么处理？

1. 用定点批次大小：固定 batch_size=128。如果你的数据末尾不足 128 条，用 0 补齐（在损失函数里 mask 掉无效位置） 2. 用 2 的幂次：batch_size 用 32, 64, 128, 256——图缓存复用率高 3. 用 padding + mask：这几乎是所有 ML 框架的标准做法——而不是"动态形状"

---

3.6 错误处理：panic 而非 error

在图定义函数里，你不能 return nil, err，因为返回类型是 *graph.Node 而非 (*graph.Node, error)。

GoMLX 的答案是：panic。

func myModel(x *graph.Node) *graph.Node {
    // 如果形状断言失败，直接 panic
    x.AssertDims(shapes.Make(dtypes.Float32, 128)) // panic 如果形状不对
    return layers.Dense(x, 64)
}

但这不会让你的程序真的崩溃。*graph.Exec 在内部捕获（recover）了 panic，转换成 error 返回：

// Must 版本：panic 不被捕获，程序崩溃
exec := graph.MustNewExec(backend, badGraph) // → panic!

// 标准版本：panic 被捕获，转成 error
exec, err := graph.NewExec(backend, badGraph)
if err != nil {
    fmt.Printf("graph build failed: %v\n", err)
    // 可以优雅降级或重试
}

为什么这么设计？

API 简洁性是首要理由。图定义函数到处都是嵌套的函数调用——如果每个都返回 error，代码会变成不可读的 error 检查地狱。panic 让图定义代码像正常 Go 代码一样干净。

另一个理由是编译期 vs 运行时的区分。图构建阶段是一个"不受信任"的阶段——你的代码可能有 bug。panic 提供了清晰的"这里出错了"的信号，而不是静默返回零值让后续运算全盘错误。

教学和原型代码用 Must* 系列（panic 友好）。生产代码用返回 error 的版本 + recover。

---

3.7 GoMLX 心智模型：一张图收束全章

┌──────────────────────────────────────────────────────────┐
│                    你的 Go 代码（图定义）                    │
│  func myFn(x *Node) *Node {                             │
│      a := graph.Square(x)    // 不计算，只画图              │
│      b := graph.MulScalar(a, 2.0)                       │
│      return b                                            │
│  }                                                        │
│                                                           │
│              ↓ 进入 graph.MustNewExec                     │
│                                                           │
│  ┌──────────────────────────────────────────────┐        │
│  │           计算图（DAG，Abstract）              │        │
│  │  [x] → [Square] → [a] → [MulScalar 2] → [b] │        │
│  └──────────────────────────────────────────────┘        │
│              ↓ 后端编译                                    │
│  ┌──────────────────────┐  ┌──────────────────────────┐  │
│  │   XLA 后端（Linux）   │  │ SimpleGo 后端（全平台）    │  │
│  │  → HLO → LLVM IR →  │  │ → Go 函数调用链 + SIMD   │  │
│  │  → CUDA/CPU 机器码   │  │ → WASM（可选）            │  │
│  └──────────────────────┘  └──────────────────────────┘  │
│              ↓ MustExec(args...)                          │
│  ┌──────────────────────────────────────────────┐        │
│  │           设备执行（Device Execution）         │        │
│  │  CPU/GPU/TPU 上真正跑数学运算                 │        │
│  └──────────────────────────────────────────────┘        │
│              ↓ 结果                                      │
│         *tensors.Tensor（搬回 Go 一侧）                   │
└──────────────────────────────────────────────────────────┘

---

3.8 在进入原理篇之前

你现在从"知道怎么写 GoMLX 代码"到达了"知道 GoMLX 怎么运作的"。这条理解链是：

> 节点（Node）→ 图（DAG）→ 编译（backend）→ 执行（Exec）

第二部分（第 4-7 章）会把这条链上的每一环拆开给你看：

  • 第 4 章：把"计算图"从抽象概念变成你用手能画的数据结构
  • 第 5 章：XLA 和 JIT 编译——从图到机器码的全过程
  • 第 6 章：张量——它的内存模型、本地 vs 设备、为什么 Finalize 不能忘
  • 第 7 章：自动微分——梯度到底是怎么精确算出来的
这些不是"高级话题"。这就是 GoMLX 的引擎。一个不理解引擎的司机可以开车，但一个理解引擎的司机才能改车、调车、在极端条件下不翻车。

翻到下一页，你拆开它。

---

> 作者：步子哥，资深AI研究员 | 连载中：《GoMLX从入门到精通》 > 上一章：第2章 GoMLX的第一行代码 > 下一章预告：第4章 计算图核心数据结构——把DAG从抽象概念拆成你用手能画的数据结构

暂无表态
💬 讨论回复 (0)
💡 推荐
话题推荐 Translational Generative Retrieval via Potential Query Generation, ICASSP 2025 话题推荐 Generative Retrieval and Alignment Model: A New Paradigm for E-commerce Retrieval, apr 2025, arxiv 话题推荐 Stealthy Attack on Large Language Model based Recommendation, Feb 2024, arxiv 话题推荐 Open Deep Research from LangChain 话题推荐 当会议室的门关上
推荐

🌟 智谱 GLM-5 已上线

我正在智谱大模型开放平台 BigModel.cn 上打造 AI 应用，智谱新一代旗舰模型 GLM-5 已上线，在推理、代码、智能体综合能力达到开源模型 SOTA 水平。

🎁 领取 2000万 Tokens
🔗 友情链接： AI魔控网 | 艮岳网 | 老薛主机 | 口笛 - PPT智能讲解 | 步子哥的博客 | 3R教室