Rust in Everything
基于 candle 0.11.0

引用与不可变

主线位置：张量主线的「计算」一步：运算怎么拿输入，结果和输入是什么关系，以及 PyTorch 的原地修改在 candle 里换成什么。
版本：PyTorch 2.15.0（CPU）、candle 0.11.0（CPU，accelerate）
flowchart TB
  t["Tensor 即 Arc"] --> cl["clone 计数加一"]
  t --> cp["copy 复制数据"]
  cl --> sh["共享 storage"]
  sh --> wr{"原地写"}
  wr -->|"slice_set 或 Var set"| both["共享者都看见"]
  wr -->|普通运算| fresh["结果是新张量"]
  cp --> own["自己的 storage"]

实验代码是一个 candle 程序 own（release，带一个只计数的全局分配器）、几个故意编不过的小程序（cargo build --color never），PyTorch 侧是 own.py；平台 M2 Max、macOS 27、rustc 1.99.0。编译器输出里的本地路径前缀换成了 ...。

运算借用，不拿走

a = torch.ones(3)
b = torch.ones(3)
c = a + b

let a = Tensor::ones(3, DType::F32, &dev)?;
let b = Tensor::ones(3, DType::F32, &dev)?;
let c = (&a + &b)?;

# PyTorch
after c = a + b: a = [1.0, 1.0, 1.0] b = [1.0, 1.0, 1.0] c = [2.0, 2.0, 2.0]
# candle
after c = &a + &b: a = [1.0, 1.0, 1.0], b = [1.0, 1.0, 1.0], c = [2.0, 2.0, 2.0]

和 PyTorch 一样，a、b 之后照常可用。candle 的 Tensor 方法全部借用：tensor.rs 里直接写出来的公开方法，122 个拿 &self，另外 17 个是 new、zeros、from_vec 这样的构造函数，没有一个按值拿 self；宏生成的运算（add、relu、broadcast_mul……）也是 &self。整个 candle-core 里，张量类型上唯一按值拿走 self 的公开方法是 Var::into_inner。

运算符则两种都实现了：bin_trait! 给 Tensor 和 &Tensor 都实现了 Add。写成 a + b（不加 &）也能编过，但这时 a、b 被移进了运算符：

let c = (a + b)?;
println!("{c} {a}");

error[E0382]: borrow of moved value: `a`
  --> src/bin/moved.rs:7:18
   |
 4 |   let a = Tensor::ones(3, DType::F32, &Device::Cpu)?;
   |       - move occurs because `a` has type `candle_core::Tensor`, which does not implement the `Copy` trait
 5 |   let b = Tensor::ones(3, DType::F32, &Device::Cpu)?;
 6 |   let c = (a + b)?;
   |           ------- `a` moved due to usage in operator
 7 |   println!("{c} {a}");
   |                  ^ value borrowed here after move
   |
note: calling this operator moves the left-hand side
...
help: consider cloning the value if the performance cost is acceptable
   |
 6 |   let c = (a.clone() + b)?;
   |             ++++++++

rustc 建议 clone()，在 candle 里更直接的改法是加 &：按值的实现里面也只是 Tensor::add(&self, rhs.borrow())，拿走所有权没有任何好处。

这一点和 burn 相反。burn 的运算按值消费张量，&a + &b 没有实现（E0369），复用要 clone()，forward(&self) 里用字段也要 self.weight.clone()，见 burn 卷的运算消费张量。candle 里同样的模块直接借用字段：

struct Layer {
  w: Tensor,
}

impl Layer {
  fn forward(&self, x: &Tensor) -> Result<Tensor> {
    x.matmul(&self.w)
  }
}

forward twice with &self.w: [[3.0, 3.0]] [[3.0, 3.0]]

clone 不复制数据

Tensor 是 Arc 包着的 Tensor_，数据在里面又一层 Arc<RwLock<Storage>>：

#[derive(Clone)]
pub struct Tensor(Arc<Tensor_>);

pub struct Tensor_ {
    id: TensorId,
    // ...
    storage: Arc<RwLock<Storage>>,
    layout: Layout,
    op: BackpropOp,
    is_variable: bool,
    dtype: DType,
    device: Device,
}

clone() 只是外层计数加一。copy() 才复制 storage，注释写着「Compared to clone, this copies the actual storage」。一个 4096×4096 的 f32 张量（64 MiB）：

y = x
z = x.clone()

let (y, mib_clone) = measure(|| x.clone());
let (z, mib_copy) = measure(|| x.copy());

# PyTorch（峰值 RSS 的增量）
x 64 MiB: y = x shares storage True, peak RSS +0 MiB
x 64 MiB: x.clone() shares storage False, peak RSS +64 MiB
# candle（分配器统计的字节数）
x 64 MiB: clone allocated 0.000 MiB, shares storage true
x 64 MiB: copy allocated 64.000 MiB, shares storage false
ids: x TensorId(4), clone TensorId(4), copy TensorId(5)

candle 的 clone() 连 TensorId 都不变，它就是同一个张量（GradStore 按 id 取梯度，所以用 clone 去取也取得到）。两边名字一样、意思不一样：

想要的效果PyTorchcandle
同一个张量再起一个名字y = xx.clone()，或者直接借用 &x
一份独立的数据x.clone()x.copy()
共享数据、断开计算图x.detach()x.detach()（共享 storage，新的 id）

把 PyTorch 的 x.clone() 照搬成 candle 的 x.clone()，只要之后没有原地写（下一节），结果没有区别：普通运算永远生成新的 storage，共享的 storage 只被读。Metal 上 copy() 有一个上游问题，它并不真正复制缓冲区，见源码卷的改了 copy 出来的张量，原张量也变了。

没有 add_，只有三个原地写的入口

PyTorch 里 t[0] = 9、t.add_(1)、t += 1 都直接改数据，指向同一个张量的名字都看得见。candle 没有 add_，也没有实现 AddAssign：

let mut a = Tensor::ones(3, DType::F32, &Device::Cpu)?;
a += &b;

error[E0368]: binary assignment operation `+=` cannot be applied to type `candle_core::Tensor`
  --> src/bin/addassign.rs:6:3
   |
 6 |   a += &b;
   |   -^^^^^^
   |   |
   |   cannot use `+=` on type `candle_core::Tensor`
   |
note: `candle_core::Tensor` does not implement `AddAssign<&candle_core::Tensor>`
  --> .../candle-core-0.11.0/src/tensor.rs:68:1
   |
68 | pub struct Tensor(Arc<Tensor_>);
   | ^^^^^^^^^^^^^^^^^ `candle_core::Tensor` is defined in another crate

写成 let a = (&a + &b)?;：新张量，旧名字被遮蔽。Tensor_ 的注释说得很明白，张量的内容可以被内部修改（inner mutability），所以 storage 外面套了 RwLock；但能写的入口只有三个：

  • slice_set：把 src 写进 self 的某一段，KV 缓存就是这么追加的。
  • Var::set：整个替换一个变量的内容，优化器用它更新参数。
  • inplace_op1/2/3：自定义原地算子，见源码卷的 inplace_op2 卡死。

它们都拿 &self，写进去的数据，所有共享这块 storage 的张量都看得见：

t = torch.tensor([1.0, 2, 3, 4])
alias = t
copied = t.clone()
t[0] = 9
m = torch.zeros(2, 3)
m[1].copy_(torch.tensor([7.0, 8, 9]))

let t = Tensor::new(&[1f32, 2., 3., 4.], &dev)?;
let alias = t.clone();
let copied = t.copy()?;
t.slice_set(&Tensor::new(&[9f32], &dev)?, 0, 0)?;
let m = Tensor::zeros((2, 3), DType::F32, &dev)?;
let row1 = m.narrow(0, 1, 1)?;
row1.slice_set(&Tensor::new(&[[7f32, 8., 9.]], &dev)?, 1, 0)?;

# PyTorch
t[0] = 9: t = [9.0, 2.0, 3.0, 4.0] alias = [9.0, 2.0, 3.0, 4.0] clone = [1.0, 2.0, 3.0, 4.0]
m[1].copy_(..): m = [[0.0, 0.0, 0.0], [7.0, 8.0, 9.0]]
# candle
t.slice_set([9], 0, 0): t = [9.0, 2.0, 3.0, 4.0], clone = [9.0, 2.0, 3.0, 4.0], copy = [1.0, 2.0, 3.0, 4.0]
m.narrow(0, 1, 1).slice_set(..): m = [[0.0, 0.0, 0.0], [7.0, 8.0, 9.0]]

这里 candle 的 clone() 表现得正是 PyTorch 的 alias = t，copy() 才是 PyTorch 的 t.clone()。经 narrow 出来的视图写，原矩阵跟着变，也和 PyTorch 的 m[1] 一样。

slice_set 有两道 PyTorch 没有的检查：

t.slice_set(&t.clone()): Err("cannot use slice_set when self and src share their storage")
non-contiguous view slice_set: Err("slice-set only supports contiguous tensors")

第二行是对 m.narrow(1, 0, 1)（一列，不连续）调 slice_set。PyTorch 的 m[:, 0] = .. 对不连续的视图照样能写。

📝NOTE
冷知识：slice_set(&self_clone) 以前会卡死

上面第一条报错来自 PR #2733「Explicit error when slice-set is called with the same src and dst.」（首次发布于 0.8.3）。之前 slice_set 在检查完连续性后直接往下走：对 self 拿写锁，再对 src 拿读锁，两者是同一把 RwLock 时永远等不到。这个 PR 在连续性检查后面加了三行 if self.same_storage(src) { crate::bail!(..) }，测试里加的注释是「This used to create a deadlock rather than returning an actual error.」，断言 cache.slice_set(&cache, 0, 0).is_err()。Var::set 开头的同一种检查（cannot set a variable to a tensor that is derived from its value）早就有了；自定义的 inplace_op2 在 0.11.0 里仍然没有，见 inplace_op2 卡死。

Var::set 对应 PyTorch 里 with torch.no_grad(): p.copy_(new) 这类参数更新。candle 的 SGD::step 就一行：

                var.set(&var.sub(&(grad * self.learning_rate)?)?)?;

var.sub(..) 先算出一个新张量，再整块写回 var 的 storage。训练怎么用 Var 和 VarMap，见训练一步。写回去以后：

w = torch.nn.Parameter(torch.tensor([1.0, 2, 3]))
seen = w.data
before = w * 2
with torch.no_grad():
  w.copy_(torch.zeros(3))

let w = Var::new(&[1f32, 2., 3.], &dev)?;
let seen = w.as_tensor().clone();
let detached = w.as_detached_tensor();
let before = (&*w * 2.0)?; // an op result is a new tensor
w.set(&Tensor::new(&[0f32, 0., 0.], &dev)?)?;

# PyTorch
after w.copy_(zeros): w.data taken before = [0.0, 0.0, 0.0] , w*2 computed before = [2.0, 4.0, 6.0]
# candle
after w.set([0,0,0]): clone = [0.0, 0.0, 0.0], w*2 computed before = [2.0, 4.0, 6.0]
w.set(&(w * 2)): Ok(())
w.set(&w.clone()): Err("cannot set variable cannot set a variable to a tensor that is derived from its value")

两边一样：拿着同一块数据的名字（w.data、as_tensor().clone()）看到新值，之前算出来的结果（w * 2）不变。w.set(&(w * 2)) 没问题，因为 w * 2 是新 storage；w.set(&w.clone()) 被拒绝，报错里 cannot set variable 是变体 CannotSetVar 自己的前缀，后面才是消息。

❗IMPORTANT
动脑筋：detach 出来的张量

上面的 detached 是 Var::set 之前用 as_detached_tensor() 取的，它断开了计算图。w.set([0, 0, 0]) 之后，detached 是 [1, 2, 3] 还是 [0, 0, 0]？

动脑筋答案

after w.set([0,0,0]): detached = [0.0, 0.0, 0.0]

是 [0, 0, 0]。as_detached_tensor 调的是 Tensor::detach，它建一个新的 Tensor_：新的 id、op 为空、is_variable: false，但 storage: self.storage.clone()，克隆的是 storage 的 Arc，数据还是同一块。断开的只是计算图，不是数据。PyTorch 的 detach() 也共享存储，w.detach() 在 w.copy_(..) 之后同样看到新值。想要一份不受影响的初始权重，用 copy()（在 Metal 上要注意上面提到的上游问题）。

小结

  • candle 的 Tensor 方法全部拿 &self，写 &a + &b，两边之后都能用；写成 a + b 会把两个都移走（E0382）。burn 正相反，运算按值消费。
  • clone() 只加引用计数、TensorId 不变，等于 PyTorch 的 y = x；要独立的数据用 copy()，那才是 PyTorch 的 x.clone()。
  • 张量没有 add_、+=（E0368）；原地写只有 slice_set、Var::set、inplace_op*，写进去的值所有共享 storage 的张量（clone、视图、detach 出来的）都看得见。
  • slice_set 要求两边连续、不共享 storage；Var::set 拒绝从自己派生出的张量，优化器的更新就是 var.set(&new)。

💬 关联讨论

发起讨论
加载中...