burn 源码解析

全卷只对着一组版本写，也就是一个新项目写 burn = "=0.22.1" 时 Cargo 解析出来的那一组：

burn 工作区写的是 cubecl = "0.11.1"、cubek = "0.3.1"；cubecl 和 cubek 是两个单独的仓库，按实验项目 Cargo.lock 里实际解析到的版本锁定。

burn 给自己的描述是「both a tensor library and a deep learning framework, optimized for numerical computing, training and inference」：用一套 Rust API 写张量运算和模型，同一份代码既用来训练也用来部署，跑在 CUDA、ROCm、Metal、Vulkan、WebGPU、CPU 和 WebAssembly 上。0.22 是一次大改：Tensor<B, D, K> 去掉了后端类型参数变成 Tensor<D, K>，后端由张量所在的 Device 在运行期决定，所有调用经过一个全局后端 Dispatch。网上和 burn 自己文档里的不少例子还是旧写法，本卷只讲 0.22.1 的样子。

本卷讲张量 API、后端抽象与分派、纯 Rust 的 CPU 后端 Flex、Autodiff、burn-ir 与 Fusion、burn-cubecl 和 cubecl 运行时、Module / Config 派生、burn-nn 的代表模块、优化器、模型存取和 Learner 主循环。GPU 实测平台是 macOS 上的 Metal（经 wgpu）；CUDA、ROCm、Vulkan、WebGPU 只读源码，凡是只读了源码的地方都会写明。cubecl 的 CPU 运行时（cpu feature，要联网下载 LLVM）、burn-remote / burn-router / burn-capture、burn-vision / burn-signal / burn-linalg / burn-rl、弃用的 burn-ndarray 和 burn-tch、量化和训练界面的渲染不展开，只在源码地图里一句话带过。

全卷沿三条线展开：

• 算子主线（一次 a.matmul(b)）：泛型的 Tensor::matmul 只做检查，转给不泛型、只经手 BridgeTensor 的 Numeric::matmul；Dispatch::float_matmul 由 #[backend_dispatch] 展开成按设备的 match；设备是 .autodiff() 的先进 Autodiff，登记一个反向步骤；GPU 上内层是 Fusion<CubeBackend>，它只把调用记成 burn-ir 的 OperationIr 送进流，立刻返回；等到 into_data 或 sync，流在 cubecl 的设备线程上被 drain，能融合的编成融合内核，matmul 交给 CubeBackend，autotune 在 cubek 的 matmul 内核里挑最快的，cubecl 把 #[cube] 内核编成 MSL 提交，结果异步读回。CPU 上同一个调用落到 Flex::float_matmul，由 gemm crate 当场算完。
• 求导主线（一次 loss.backward()）：在 Device::autodiff() 上建的张量带着图节点；每个可求导的 op 登记一个实现 Backward 的步骤和它需要的状态（或者按检查点策略只记怎么重算）；backward() 从根节点收集步骤、按顺序执行，梯度按节点 id 累加进 Gradients；GradientsParams 再把它们按参数 id 重新组织给优化器。
• 训练主线（一个 step）：#[derive(Module)] 生成 visit / map 遍历；forward 是一串算子主线；loss.backward() 是一次求导主线；ModuleOptimizer::step 用 map 遍历模块，按参数 id 取梯度和 Adam 状态、换上新参数；Learner 把循环、数据加载、指标和 checkpoint 包起来。

三条线每一步对应的源码和页面，见怎么读这一卷。

flowchart TB
  intro[00 导读] --> overview[01 全景]
  overview --> tensor[02 张量]
  tensor --> backend["03 后端与分派"]
  backend --> ad[04 自动微分]
  backend --> fusion["05 IR 与融合"]
  fusion --> cubecl["06 CubeCL 与 GPU"]
  ad --> module["07 模块与训练"]
  cubecl --> design[08 设计取舍]
  module --> design
  design --> pit[09 常见坑]
  pit --> site[10 在本站中]
  site --> hands[11 动手]

实验环境

关于行为的说法（一次调用经过哪些层、在哪个线程上提交内核、提交了几个内核、首次和之后各花多久、rustc 和 panic 报什么、宏展开成什么）都配一个小程序和它的真实输出，平台是 M2 Max、macOS 27、rustc 1.99.0。实验依赖锁定 burn = "=0.22.1"，一律离线构建；每个实验写明开了哪些 feature、是 debug 还是 release。GPU 计时都在结束处 device.sync()，并把第一次（编译内核、autotune）和之后分开；cubecl 会把 autotune 结果缓存在最外层 Cargo.toml 旁的 target/environment/ 里，所以「新进程的第一次」不一定是冷的，每处都写明缓存状态。训练实验只用合成数据，不下载数据集。

各部分

1. 00 导读：锁定的版本，0.22 的大改，三条主线每一步落在哪个函数、哪一页讲；feature 怎么在 burn-dispatch 里变成三个 cfg；工作区按层分组。
2. 01 全景：burn 和 candle、tch-rs、PyTorch 的分界；一次 matmul、一个训练步各从头到尾走一遍，带调用栈、线程名和计时。
3. 02 张量：Tensor<D, K>、种类 trait、BridgeTensor 和 *_impl、Device 门面、形状检查与报错。
4. 03 后端与分派：Backend 这组 trait，Dispatch 怎么按设备 match，纯 Rust 的 CPU 后端 Flex。
5. 04 自动微分：Autodiff 装饰器，一个算子的反向，backward 与 Gradients，梯度检查点。
6. 05 IR 与融合：burn-ir，Fusion 装饰器，融合搜索，生成融合内核。
7. 06 CubeCL 与 GPU：#[cube] 内核，运行时与 Client，wgpu 与 Metal，autotune，burn-cubecl 后端。
8. 07 模块与训练：Module 派生、Config、nn 模块、优化器、存储与加载、Learner 与数据。
9. 08 设计取舍：后端是类型参数还是运行期分派，装饰器套娃，惰性加融合，自己写 GPU 语言，和 candle 对照。
10. 09 常见坑：用报错原文或现象做标题，每个都先跑出最小复现。
11. 10 在本站中：本站不用 burn；给本站加一个「相关文章」功能、用 burn 来算，代价是什么。
12. 11 动手：不依赖 burn，约 300 行写一个反向模式自动微分，再和 burn 的实现逐项对照。

适合谁读

• 用 burn 写过模型，想知道 matmul 在 GPU 上为什么「返回得那么快」、第一次运行为什么慢、backward 之后梯度怎么到了优化器手里的读者：从导读开始，按阅读路线里「只想会用」那条读。
• 想读懂一个 Rust 深度学习框架怎么分层的读者：按顺序读。03 和 04 讲清楚后端和装饰器之后，05、06 的融合与 GPU 才好读。
• 需要的前置知识：Rust 的 trait、关联类型和 const 泛型，一点深度学习（张量、反向传播、优化器）；Config 和记录走 serde（见 serde 卷），tracing feature 打开的是 span（见 tracing 卷），GPU 读回是 async fn（future 的基本概念见 tokio 卷）。