第2章：GoMLX 的第一行代码

---

2.1 搭建环境

2.1.1 Go 版本与 GoMLX 安装

GoMLX v0.27.3 要求 Go 1.22 及以上版本。先确认你的版本：

go version
# 输出示例：go version go1.26.2 windows/amd64

如果低于 1.22，参考 go.dev/dl 升级再来。

然后创建项目并安装：

mkdir gomlx-book && cd gomlx-book
go mod init gomlx-book
go get github.com/gomlx/gomlx@v0.27.3

首次 go get 会下载所有依赖。GoMLX 的核心依赖极少：它只链接 backends/default 包（内含 XLA 或纯 Go 后端）和一些必要的工具库。整体依赖树会在 go mod tidy 后自动拉齐。

2.1.2 两类计算后端，一个统一接口

GoMLX 的设计理念之一：你的图定义代码不感知底层硬件。它通过统一的 backends.Backend 接口，让运行时选择用哪个后端执行计算。

GoMLX v0.27.3 的后端体系：

2.1.3 创建后端的正确方式

package main

import (
    "fmt"

    // 这一行注册 XLA 和 SimpleGo 后端。务必保留！
    _ "github.com/gomlx/gomlx/backends/default"

    "github.com/gomlx/gomlx/backends"
)

func main() {
    // backends.New() 自动选择可用的后端：
    //   Linux/macOS → XLA（优先）
    //   Windows     → SimpleGo（纯 Go）
    //   或通过 $GOMLX_BACKEND 环境变量手动指定
    backend, err := backends.New()
    if err != nil {
        panic(fmt.Sprintf("failed to create backend: %v", err))
    }
    fmt.Printf("Using backend: %s\n", backend.Description())
    // Windows 输出：Using backend: Simple Go Portable Backend
    // Linux   输出：Using backend: XLA (CPU)
}

两个要点：

1. _ "github.com/gomlx/gomlx/backends/default" 必须保留——这行代码通过 Go 的 init() 机制注册所有可用后端。删掉它，backends.New() 会 panic 说"没有注册的后端"。

2. backends.New() 返回的是 backends.Backend 接口——它抽象了所有后端类型。你的图定义代码只需要和这个接口打交道。

---

2.2 第一个计算图：1 + 1 = 2

Go 程序员写 1 + 1，一行结束：

result := 1 + 1  // 2。直接执行，返回 int。

GoMLX 不是这样。在 GoMLX 里你写的代码不是直接执行计算，而是构建一张计算图。图构建完了，才编译，才执行。

2.2.1 定义图——此时计算还没发生

import "github.com/gomlx/gomlx/pkg/core/graph"

// 这是一个"图定义函数"——它描述计算，但不执行计算。
// *graph.Node 代表图中的每一个运算结果。
addGraph := func(a, b *graph.Node) *graph.Node {
    return graph.Add(a, b)
}

这段代码里没有 int，没有 float64，没有任何"数值"——只有 *graph.Node 之间的运算关系。

直觉：这是"写菜谱"阶段。你定义了"输入两个食材 → 做加法 → 输出成品"，但还没开火。

2.2.2 编译图

// graph.MustNewExec 做三件事：
//   1. 调用 addGraph 构建完整的计算图（记录所有运算）
//   2. JIT 编译这张图（SimpleGo 或 XLA 后端）
//   3. 返回一个 *graph.Exec 对象，可反复调用
exec := graph.MustNewExec(backend, addGraph)

MustNewExec 内部发生了什么： 1. 调用 addGraph——此刻它传入两个"占位"的 *Node，记录下"这里有一次 Add 运算" 2. 把整张图交给后端编译 3. 返回 *graph.Exec——这是一个编译后可以反复调用的函数

Must 前缀：如果编译失败，直接 panic。生产代码可改用 graph.NewExec（返回 error）。教学代码用 Must 系列更干净。

2.2.3 执行——这时加法才真正发生

// exec.MustExec 传入实际的 Go 数值
// 返回 []*tensors.Tensor 切片
results := exec.MustExec(1.0, 2.0)
result := results[0]

// result.Value() 返回 any，对标量而言就是 Go 原生数值
fmt.Printf("1.0 + 2.0 = %v\n", result.Value())
// 输出：1.0 + 2.0 = 3

关键的认知时刻：你传给 MustExec 的 1.0 和 2.0 只是普通的 Go 字面量。GoMLX 把它们搬到了设备内存（如果有 GPU 就在显存），跑完加法，把结果搬回 Go 可读的内存。

整个流程总结：

---

2.3 手动构图 vs 函数式构图

GoMLX 支持两种写计算图的方式。上面演示的是函数式构图——定义一个普通函数，GoMLX 自动追踪运算。这也是本书的默认方式。

还有一种手动构图，在需要命名节点或动态修改图结构时有用：

import (
    "github.com/gomlx/gomlx/pkg/core/dtypes"
    "github.com/gomlx/gomlx/pkg/core/shapes"
)

// 手动创建一张空图
g := graph.NewGraph(backend, "my-add-graph")

// 手动创建参数节点
a := graph.Parameter(g, "a", shapes.Make(dtypes.Float64))
b := graph.Parameter(g, "b", shapes.Make(dtypes.Float64))

// 手动添加运算
c := graph.Add(a, b)

// 手动编译——指定 c 为输出
exec := g.Compile(c)

// 执行
results := exec.MustExec(1.0, 2.0)
fmt.Printf("%v\n", results[0].Value()) // 3

对比：

2.4 第一个"真实"例子：f(x) = x² + 2x + 1

1 + 1 = 2 太简单了。上一个稍微显露肌肉的例子——多项式函数，用不同输入反复调用同一个编译好的图。

package main

import (
    "fmt"

    _ "github.com/gomlx/gomlx/backends/default"
    "github.com/gomlx/gomlx/backends"
    "github.com/gomlx/gomlx/pkg/core/dtypes"
    "github.com/gomlx/gomlx/pkg/core/graph"
)

func main() {
    backend, _ := backends.New()
    fmt.Printf("Backend: %s\n\n", backend.Description())

    // 定义图：f(x) = x² + 2x + 1
    // graph.Square(x) → x²
    // graph.MulScalar(x, 2.0) → 2x
    // graph.Scalar(g, dtypes.Float64, 1.0) → 常量 1.0
    polynomial := func(x *graph.Node) *graph.Node {
        xSquared := graph.Square(x)
        twoX := graph.MulScalar(x, 2.0)
        one := graph.Scalar(x.Graph(), dtypes.Float64, 1.0)
        return graph.Add(xSquared, graph.Add(twoX, one))
    }

    // 编译——一次，永久
    exec := graph.MustNewExec(backend, polynomial)

    // 用不同 x 反复执行——每次执行都是纯粹的数学计算
    testValues := []float64{-3.0, -2.0, -1.0, 0.0, 1.0, 2.0, 3.0}
    fmt.Printf("f(x) = x² + 2x + 1\n")
    for _, x := range testValues {
        res := exec.MustExec1(x)
        fmt.Printf("  f(%.0f) = %.0f\n", x, res.Value())
    }
}

实测输出（Windows 10, Go 1.26.2, SimpleGo 后端）：

Backend: Simple Go Portable Backend

f(x) = x² + 2x + 1
  f(-3) = 4
  f(-2) = 1
  f(-1) = 0
  f(0) = 1
  f(1) = 4
  f(2) = 9
  f(3) = 16

验证：f(-3) = 9 + (-6) + 1 = 4 ✓。f(2) = 4 + 4 + 1 = 9 ✓。

2.4.1 这段代码教给我们的几样东西

第一，图编译一次，执行多次。

graph.MustNewExec 只调用一次——编译了 f(x)=x²+2x+1 的完整计算图（含常量折叠等优化）。之后七次 exec.MustExec1(x) 都是直接在编译好的机器码上跑——不重新构图、不重新编译。

第二，图构建阶段的常量是"固化的"。

one := graph.Scalar(x.Graph(), dtypes.Float64, 1.0)

graph.Scalar 创建的常量节点——值 1.0 在编译时就被固化。无论 MustExec1(100) 还是 MustExec1(-500)，这个 1.0 永远不变。只有通过 graph.Parameter 定义的节点（见 2.3 节手动构图示例），才能在每次执行时传入不同的值。

**第三，graph.Add、graph.MulScalar 返回的是 *graph.Node，不是数值。

在 polynomial 函数内部，你所做的一切都是"把节点串起来"——真正的数值计算在 MustExec1 里才启动。

---

2.5 后端切换实验

现在我们做一个小实验：在同一个机器上用同一段图定义，测后端吞吐。

Windows 上只能跑 SimpleGo，但我们仍然可以观察到它的一致性和性能特征。以下是在 Windows 10 / Go 1.26.2 / SimpleGo 后端的实测代码：

func main() {
    backend, _ := backends.New()
    fmt.Printf("Backend: %s\n", backend.Description())

    // 定义图：f(x) = x⁴ + 1
    f4 := func(x *graph.Node) *graph.Node {
        x2 := graph.Square(x)
        x4 := graph.Square(x2)
        one := graph.Scalar(x.Graph(), dtypes.Float64, 1.0)
        return graph.Add(x4, one)
    }

    exec := graph.MustNewExec(backend, f4)

    numRuns := 100000
    start := time.Now()
    for i := 0; i < numRuns; i++ {
        x := float64(i%100)*0.1 + 1.0
        res := exec.MustExec1(x)
        res.FinalizeAll() // 立即释放设备内存
    }
    elapsed := time.Since(start)

    fmt.Printf("[%s] %d runs in %v (%.0f runs/sec)\n",
        backend.Name(), numRuns, elapsed,
        float64(numRuns)/elapsed.Seconds())
}

实测输出：

Backend: Simple Go Portable Backend
[SimpleGo (go)] 100000 runs in 759.3ms (131694 runs/sec)

2.5.1 几个值得注意的结论

SimpleGo 后端已经"够快"——每秒 13 万次调用，对绝大多数 CPU 推理场景足以胜任。XLA 后端会更快（在我们未测试的 Linux 上，预计快 3-5 倍）。

FinalizeAll()——执行结果张量占用设备内存。调用 FinalizeAll() 立即释放。不释放的话，在百万次循环后会内存泄漏。第 6 章会专门讲张量的内存管理。

编译产物的大小——本节代码（含 GoMLX 的 SimpleGo 后端）编译结果：

$ go build -o myprogram.exe .
$ ls -lh myprogram.exe
-rwxr-xr-x  1 user  staff   7.8M  myprogram.exe

# 用 -ldflags="-s -w" 剥离调试信息后：
$ go build -ldflags="-s -w" -o myprogram.exe .
$ ls -lh myprogram.exe
-rwxr-xr-x  1 user  staff   5.6M  myprogram.exe

一个不带任何模型权重的"空"GoMLX 程序，编译后约 5.6-7.8MB（含 GoMLX 的运行时和 SimpleGo 后端）。当你加载实际模型权重时，体积会加上权重文件的大小——一个典型 CTR 模型（如 DeepFM，约 500 万~2000 万参数）会额外增加约 20-80MB。

---

2.6 常见踩坑与排查

坑一：no registered backends panic

现象：

panic: no registered backends for GoMLX -- maybe import the default ones?

原因：忘了导入 _ "github.com/gomlx/gomlx/backends/default"。

解决：在 import 块加上这行空白导入。它是 GoMLX 的"发动机钥匙"——插上才有后端。

坑二：graph.Scalar 编译报错

现象：写了 graph.Scalar(g, 1.0) 然后编译不过。

原因：GoMLX 的 graph.Scalar 需要显式指定数据类型。

解决：必须传 dtypes.Float64（或 dtypes.Float32 等）：

// ❌ 错误
one := graph.Scalar(g, 1.0)

// ✓ 正确
import "github.com/gomlx/gomlx/pkg/core/dtypes"
one := graph.Scalar(g, dtypes.Float64, 1.0)

坑三：MustExec 返回 shape mismatch

现象：编译正常，但运行 MustExec 时报 shape 不匹配。

原因：图的输入形状在编译时就被锁定。比如第一次调用 MustExec(1.0, 2.0)（两个标量），第二次调用 MustExec([]float64{1.0, 2.0})（一个向量），形状不同，编译期就无法匹配。

解决**：要么统一输入形状，要么让 GoMLX 缓存不同形状的图—— *graph.Exec 默认缓存最多 32 张不同形状的图，足够日常使用。

坑四：依赖下载慢

现象：go get 或 go mod tidy 卡在 github.com/gomlx/go-xla。

原因：XLA 相关包较大（~200MB 静态库），且可能需科学上网。

解决：

• 如果你只需要在 Windows/macOS 上用 SimpleGo：XLA 包不会被编译进去（有 build tag 限制），只影响下载速度
• 设置 Go 模块代理：GOPROXY=https://goproxy.cn,direct go mod tidy

2.7 本章小节，与下一章的桥

你现在手上有了：

1. 一个能跑的 GoMLX 环境（Windows 上是 SimpleGo 后端，Linux 是 XLA） 2. 一个完整的两阶段执行流——定义图 → 编译图 → 执行图 3. 三次实测：1+1=2、多项式求值、吞吐量基准（约 13 万次/秒）

如果你现在还残留一个疑问——"为什么图定义时不直接执行？为什么要先建图再编译？"——那太好了。这正是第 3 章要回答的核心问题。

第 3 章会深入 GoMLX 的世界观：计算图是什么数据结构、JIT 编译做了什么、以及为什么这种设计给了 GoMLX 跨后端、跨加速器的能力。

你不需要完全理解它才能用 GoMLX——就像你不需要理解 LLVM 才能用 Go。但理解了它，你就不再是"调 API 的人"，而是"知道引擎怎么转的人"。

打开 Go 编辑器，把上面的代码敲进去——跑通了再翻页。

---

> 作者：步子哥，资深AI研究员 | 连载中：《GoMLX从入门到精通》 > 上一章：第1章 机器学习是什么——一个Go工程师的视角 > 下一章预告：第3章 GoMLX的世界观——计算图、后端与执行模型