一次 grad

主线位置：traced 主线从头到尾，最后用 eager 主线求同一个梯度对照。本页用最小的例子 y = Σ xᵢ²（梯度是 2x）把每一步能看到什么、花多久量出来；每一步的机制在 04、05 各页细讲。

sequenceDiagram
  participant U as 用户代码
  participant T as TracedTensor
  participant A as tenferro-ad
  participant C as GraphCompiler
  participant R as Runtime
  participant E as CpuEngine
  U->>T: from_vec_col_major
  U->>T: x 乘 x 再 reduce_sum
  T-->>U: 只记图 返回 Result
  U->>A: y.grad x
  A-->>U: 新的 TracedTensor
  U->>C: compile
  C-->>U: CompiledGraph
  U->>R: builder 注册 engine 描述
  U->>R: prepare 按 engine.target
  U->>R: run 借出 engine
  R->>E: 逐条执行
  R-->>U: Vec Output
  U->>E: into_tensor

程序

和 README 的 Quickstart B 一样，只是把各步拆开计时（one_grad，release）：

use tenferro_ad::TracedTensorAdExt;
use tenferro_cpu::CpuEngine;
use tenferro_runtime::{ExecutionEngine, GraphCompiler, Runtime, TracedTensor};

let x = TracedTensor::from_vec_col_major(vec![n], data.clone())?;
let y = (&x * &x)?.reduce_sum(Some(&[0]))?;
let g = y.grad(&x)?;

let mut compiler = GraphCompiler::new();
let program = compiler.compile(&g)?;

let mut engine = CpuEngine::serial();
let mut builder = Runtime::builder();
builder.register_engine(tenferro_cpu::runtime_engine_registration(&engine)?)?;
let runtime = builder.build()?;
let prepared = runtime.prepare(&program, &[], &[engine.target()])?;
let mut out = runtime.run(&mut [&mut engine], &prepared, &[])?;
let grad = out.remove(0).into_tensor(&mut engine)?;

README 的例子把编译、建 engine、建 Runtime、prepare 都放在一个 run 函数里，每求一个值就全来一遍；这里拆开，是为了看清哪些步骤可以只做一次。target() 是 ExecutionEngine trait 的方法，README 的例子也 use 了这个 trait。

第一步：运算只记图

TracedTensor 不装数据，装的是图里的一个位置：

pub struct TracedTensor {
    pub id: TracedTensorId,
    pub rank: usize,
    pub dtype: DType,
    pub(crate) graph: Arc<Graph<StdTensorOp>>,
    pub val: LocalValueId,
    pub(crate) data: Option<Arc<RetainedValue>>,
    pub(crate) shape_hint: Option<Vec<SymDim>>,
    // ...
}

（TracedTensor，后面还有输入表、叶子元数据、检查点链等字段。）from_vec_col_major 建的输入把数据放在 data 里，执行时作为默认输入，所以 prepare 和 run 的输入都传 &[]。每个运算建一个新的 computegraph Graph 片段，用 parents 指回输入所在的片段。沿 parents 把 y 背后的片段全部走一遍：

fn graph_size(t: &TracedTensor) -> (usize, usize) {
  let mut seen = HashSet::new();
  let mut stack = vec![Arc::clone(t.graph())];
  let mut ops = 0;
  while let Some(g) = stack.pop() {
    if !seen.insert(Arc::as_ptr(&g) as usize) {
      continue;
    }
    ops += g.operations().len();
    for p in g.parents() {
      stack.push(Arc::clone(p));
    }
  }
  (seen.len(), ops)
}

y graph (fragments, ops): (3, 2)
-- y traced graph (own fragment)
  ReduceSum { axes: [0] } role=Primary inputs=1

三个片段（输入 x、Mul、ReduceSum）、两个运算，什么都还没算。形状和 dtype 在记图时就检查，所以 &x * &x 的类型是 Result<TracedTensor, Error>，链式写法每步都要 ?（见怎么读这一卷）。图怎么建、为什么是一个个片段，见 computegraph 和 TracedTensor。

第二步：grad 生成一张新图

TracedTensorAdExt::grad 来自 tenferro-ad，tenferro-runtime 自己不会求导。它的实现取默认的 semantic 扩展规则集，走 grad_with_optional_rules：

impl TracedTensorAdExt for TracedTensor {
    fn grad(&self, wrt: &TracedTensor) -> Result<TracedTensor> {
        let rules = SemanticExtensionRuleSet::default();
        grad_with_optional_rules(self, wrt, &rules, None)
    }
    // ...

返回的还是一个 TracedTensor，背后是一张新图：

grad graph (fragments, ops): (1, 4)
-- grad traced graph (own fragment)
  BroadcastInDim { shape: [Const(3)], dims: [] } role=Primary inputs=1
  Mul role=Primary inputs=2
  Mul role=Primary inputs=2
  Add role=Primary inputs=2

只有一个片段，四个运算，正向的 Mul 和 ReduceSum 都不在里面：ReduceSum 的反向是把标量种子（全 1）广播回 [3]，x * x 的反向是「种子乘另一边」两项相加，两项都用不到正向的结果。形状 [Const(3)] 是具体的数：用 from_vec_col_major 建的输入形状已知，形状直接写进了图。规则放在哪、grad 怎么由线性化和转置组成，见 primitive AD 和 traced AD。

第三步：编译

GraphCompiler::compile 把各个片段解析、交给 computegraph 的 materialize_merge（调用处）压平成一张图，再冻结成 CompiledGraph。它的 Debug 只给统计：

CompiledGraph { inputs: 2, outputs: 1, input_keys: 2, bindings: 2, semantic_fingerprint: SemanticFingerprint([145, 199, 23, 76, 188, 78, 19, 74, 36, 185, 105, 37, 156, 98, 92, 20]) }

两个输入、两个绑定：x 的数据，和求导时造出来的那个全 1 种子。CompiledGraph 里没有任何 engine 的东西，program() 给出的 SemanticProgram 是一串与后端无关的 SSA 运算；program().operations() 可以逐个取出，SemanticOpRef::Core 里装的是 CoreSemanticOp，match 出来就能打印。

指纹里有形状：同一个程序，n = 3 时是 [145, 199, 23, ...]，n = 1000000 时是 [17, 151, 37, ...]。编译只要十几到二十微秒（下表），同一个 GraphCompiler 再编一次、换一个新的 GraphCompiler 编，时间都差不多。编译做了哪些事，见编译。

第四步：Runtime、prepare 和 run

Runtime 不认识 CpuEngine 这个类型，只认识 EngineRegistration：tenferro_cpu::runtime_engine_registration 把一个 engine 翻译成注册项，register_engine 收下，build 出一个不可变的 Runtime。注册项只是描述，Runtime 里不存 engine。

执行分两步。prepare 收程序、输入签名和一组目标描述（engine.target()），给出一个 Prepared；run 收 &mut [&mut dyn ExecutionEngine]，第 i 个槽位就在第 i 个 engine 上跑，返回 Vec<Output>。CPU 上 run 返回时活已经干完；Output 要么是新的紧凑张量，要么是某个输入的借用视图，into_tensor 统一转成 Tensor。结果是 grad[..3] = [2.0, 4.0, 6.0]，y 是 [14.0]。上游的设计记录 execution-engine-ownership.md 把这三步的频率写成一张表：compile「once per computation」，prepare「once per input signature and target set」，run「as often as needed」。

同一个进程里第一次执行比之后慢。为了分清这笔钱付给了谁，又做了几件事（串行 engine，n = 3，一次运行）：

prepare grad                                   92.8 us
run grad (first)                               51.1 us
into_tensor                                     0.1 us
run grad (later)                                4.2 us (min of 50)
prepare+run+into_tensor (later)                23.7 us (min of 50)
prepare y (same runtime)                       33.7 us
run y (first, same runtime)                    27.3 us
run recompiled grad (same runtime)              5.4 us
warm_up then run                                5.7 us
run grad (fresh engine + runtime)               5.5 us

• 第一次 run 是 51 µs，之后 4.2 µs。同一个 Runtime 上第一次跑另一个程序 y（Mul 加 ReduceSum，ReduceSum 是这个进程里第一次出现）是 27 µs；而重新编译、重新 prepare 的 grad 程序，或者新 engine 加新 Runtime 上的 grad 程序，第一次 run 都只要五六微秒。所以这笔一次性的钱主要付给了进程里第一次用到的那几种运算，不是每个程序、每个运行时各付一次。
• prepare 比 run 贵得多：第一次 93 µs，之后每次 prepare + run + into_tensor 是 24 µs，比单独 run 多出约 20 µs。要反复执行就只 prepare 一次，把 Prepared 留着。
• Runtime::warm_up 是 0.8.0 新加的：它就是一次普通的 run，丢掉输出，再对每个 engine synchronize()，让首次执行的代价落在调用方选的时刻。CPU 上它和一次 run 差不多。prepare 缓存什么、run 每条命令付什么，见执行与缓存。

同一个梯度，eager 来一遍

use tenferro_ad::{EagerRuntime, Tensor};

let rt = EagerRuntime::new()?;
let mut engine = CpuEngine::serial();
let x = rt.variable_from(Tensor::from_vec_col_major(vec![n], data.clone())?)?;
let loss = rt.with_eager_session(&mut engine, |s| {
  let p = s.mul(&x, &x)?;
  s.reduce_sum(&p, Some(&[0]))
})?;
loss.backward(&mut engine)?;
println!("{:?}", &x.grad()?.expect("grad").as_slice::<f64>()?[..3]); // [2.0, 4.0, 6.0]

没有编译、没有 prepare：EagerRuntime::new 只建规则集和缓存，「The context stores no engine」；每个运算写在 with_eager_session 的闭包里，当场在借来的 engine 上算出值，用的是和 traced 同一套 StdTensorOp；backward 只接受标量，也要借一个 engine，梯度累加进 x 的 grad() 槽。累加是真的：backward 的文档示例里对 x + x 求两次，grad() 是 [4.0, 4.0, 4.0]。所以计时循环里每轮先 clear_grad()。eager 怎么记录、和 traced 共用哪些规则，见 eager AD。

各阶段耗时

release，负载 2.3 左右。每格是一次运行里的数（同样的参数各跑了 3 次，表里取第二次，三次的「之后」相差不到 20%）；单次的阶段是那次运行里的单次计时，「之后」是 50 次（n = 3）或 20 次（n = 10⁶）里的最小值，单位微秒。「串行」是 CpuEngine::serial()，「8 线程」是 8 个 worker 的池：

几件事：

• 记图、求导、编译都便宜，都是十几到一百多微秒，而且和 n 无关（n = 10⁶ 时 from_vec_col_major 慢，是因为要搬 8 MB 数据）。可以放心地在程序启动时建好图、编译好、prepare 好，反复执行。
• 小数据上，线程池不再拖后腿。 n = 3 时一次 run 串行 4.2 µs、8 线程 4.5 µs，eager 一轮都是 46 µs 左右；EagerRuntime::new 只要几微秒，因为它不再建线程池。设计记录说令牌只在某个入口把活拆成多份时才进池（「a small operation pays no pool entry」），n = 3 时 8 线程和串行几乎一样快，和这一点一致。
• 大数据上，多线程和 traced 都值。 n = 10⁶ 时 traced 执行串行 368 µs、8 线程 54 µs；eager 一轮是 2.5 ms 和 1.3 ms。traced 把 3 个运算作为一个准备好的程序交给 engine；eager 每轮还要立即算正向（x*x 和求和），再在记录上展开反向。差在哪里，见执行与缓存和 eager AD。
• 12 线程在本机不划算。 同样的程序用 12 线程的 engine，n = 10⁶ 的 run 之后最快 180–230 µs，比 8 线程的 54 µs 慢，单次计时还会跳到几十毫秒；原因和一次 matmul 里 8 线程比 12 线程快一样，team 要等能效核。

动脑筋答案

3 个。编译后的程序（n = 3，print_op 对 SemanticOpRef::Core 做 match 后打印）：

SemanticProgram { inputs: 2, outputs: 1, values: 5, operations: 3, shape_guards: 0 }
...
-- grad program
  BroadcastInDim { shape: [Const(3)], dims: [] } [ProgramValue(<opaque>)] -> [ProgramValue(<opaque>)]
  Mul [ProgramValue(<opaque>), ProgramValue(<opaque>)] -> [ProgramValue(<opaque>)]
  Add [ProgramValue(<opaque>), ProgramValue(<opaque>)] -> [ProgramValue(<opaque>)]

• grad 图里的两个 Mul 是 x * x 两边各自的反向项，都是「广播后的种子乘 x」。computegraph 的物化按运算的键（运算本身加上输入的键）去重，MaterializedGraph 的文档就叫它「Fully flattened, deduplicated graph」。编译后只剩一个 Mul，说明两项的键相同，被合成了一个，Add 的两个输入是同一个值。
• 剩下的是 BroadcastInDim（种子广播成 [3]）、Mul、Add：梯度 = s·x + s·x = 2x。
• values: 5 是两个输入（x、种子）加三个运算的输出。正向的 y 没有进这个程序，它在 grad 那一步就已经不在图里了。

小结

• TracedTensor 的运算只往 computegraph 的图片段里加节点，片段之间用 Arc 指向父片段；形状和 dtype 在记图时检查，所以返回 Result。
• grad（tenferro-ad）生成一张新图；y = Σx² 的梯度图是 4 个运算，不含正向；from_vec_col_major 的具体形状直接写进图。
• compile 物化、去重、冻结成与后端无关的 SemanticProgram：两个相同的 Mul 合成一个，剩 3 个运算；Debug 只给统计，看运算要 match 出 CoreSemanticOp。
• Runtime 只存注册的描述；prepare 按 engine.target() 准备一次，run 借调用方的 engine 反复执行，into_tensor 取结果。n = 3 时一次 run 约 4 µs，prepare 每次要多花约 20 µs；进程里第一次用到某种运算时多付几十微秒。
• eager 求同一个梯度不用编译，运算写在 with_eager_session(&mut engine, ..) 里，backward(&mut engine) 往 grad() 槽里累加；大数据上比 traced 慢几倍。