Tensor 的结构

主线位置：算子主线的第二步。Tensor::matmul 检查完形状，写的是 self.storage().matmul(&rhs.storage(), ..)：两次 storage() 各取一把读锁，拿到 Storage 枚举交给后端。本页讲这把锁在哪、锁住的是什么、结果张量怎么包起来；Storage 枚举怎么按设备分派见 Storage 与后端 trait。

flowchart TB
  t["Tensor"] --> a1["Arc Tensor_"]
  a1 --> id["id TensorId"]
  a1 --> lay["layout"]
  a1 --> op["op BackpropOp"]
  a1 --> meta["dtype device is_variable"]
  a1 --> a2["Arc RwLock Storage"]
  a2 --> cpu["Cpu CpuStorage"]
  a2 --> mtl["Metal MetalStorage"]
  v["视图 Tensor"] --> a1v["另一个 Tensor_"]
  a1v --> a2

两层 Arc 各管一件事：外层让张量本身可以随手复制（计算图里到处存着输入张量），内层让多个张量共用一块数据（视图）。图里下方那个「视图」张量有自己的 Tensor_，但和原张量指向同一个 RwLock<Storage>。

本页的实验 crate 是 exp-candle/tensor-exp（candle 0.11.0，features metal accelerate；M2 Max、macOS 27、rustc 1.99.0），除特别说明外都是 release。

两层 Arc

Tensor 只有一个字段：

// Tensors are refcounted so that cloning is cheap when building the op graph.
// Storages are also refcounted independently so that its possible to avoid
// copying the storage for operations that only modify the shape or stride.
#[derive(Clone)]
/// The core struct for manipulating tensors.
// ...
pub struct Tensor(Arc<Tensor_>);

它实现了 Deref<Target = Tensor_>，所以方法体里可以直接写 self.layout、self.storage。Tensor_ 有七个字段：

字段	类型	作用
`id`	`TensorId`	进程内唯一的编号，来自一个全局 `AtomicUsize` 计数器；求导时 `GradStore` 按它找梯度
`storage`	`Arc<RwLock<Storage>>`	数据本身，`Storage` 是 `Cpu` / `Cuda` / `Metal` 三个变体的枚举
`layout`	`Layout`	形状、步长、起始偏移，见布局与视图
`op`	`BackpropOp`	产生这个张量的算子和输入；不需要求导时是 `None`，见计算图怎么记下来
`is_variable`	`bool`	是否是 `Var` 背后的张量
`dtype`	`DType`	元素类型
`device`	`Device`	所在设备

dtype 和 device 在 Storage 里也能查到（Storage::dtype()、Storage::device()）。这两个字段是 storage 被放进可变容器的同一个提交里加的（见下面的冷知识），有了它们，Tensor::dtype 和 Tensor::device 直接返回字段，不用取锁。

新张量统一由 from_storage 包起来，它总是用连续布局：

pub(crate) fn from_storage<S: Into<Shape>>(
    storage: Storage,
    shape: S,
    op: BackpropOp,
    is_variable: bool,
) -> Tensor {
    let dtype = storage.dtype();
    let device = storage.device();
    let tensor_ = Tensor_ {
        id: TensorId::new(),
        storage: Arc::new(RwLock::new(storage)),
        layout: Layout::contiguous(shape),
        op,
        is_variable,
        dtype,
        device,
    };
    Tensor(Arc::new(tensor_))
}

算子主线最后一步「新 Tensor 带上 BackpropOp」就是这一句：matmul 拿到后端返回的 Storage，调 from_storage(storage, c_shape, op, false)。所以每个算子的结果都是新的 id、新的 storage、连续布局。

在 share 里打印几个类型的大小：

size_of Tensor = 8
size_of Layout = 56
size_of BackpropOp = 56
size_of Device = 80
size_of DType = 1
size_of Storage = 96
size_of Error = 80

Tensor 是一个指针。Tensor_ 不公开，不能直接 size_of，但它的字段都在上面：Layout 是两个 Vec（形状、步长）加一个偏移；Device 要装下 MetalDevice，有 80 字节。

为什么是 RwLock

Tensor_ 里 storage 字段上方的注释把候选方案列了一遍：

    // As we provide inner mutability on the tensor content, the alternatives are:
    // - Using a mutex, this would have the highest cost when retrieving the storage but would
    //   prevent errors when concurrent access takes place. Mutex would also be subject to
    //   deadlocks for example using the current code if the same tensor is used twice by a single
    //   binary op.
    // - Using a refcell unsafe cell would have some intermediary cost, borrow checking would be
    //   verified dynamically, but the resulting tensors would not be send or sync.
    // - Using an unsafe cell would have the lowest cost but undefined behavior on concurrent
    //   accesses.
    // Ideally, we would use Arc<Storage> for tensors on which we don't plan on modifying the data
    // and Arc<Mutex<Storage>> for tensors where the data could be modified, e.g. variables but
    // that's tricky to encode in the current setup.
    storage: Arc<RwLock<Storage>>,

要内部可变性，是因为有少数操作要原地改数据，最典型的是优化器更新参数：Var::set 取写锁，把新值拷进 Var 自己的 storage（见优化器与训练循环）。RwLock 在注释列的几条路里取了中间：

和 Mutex 比：a + a 这种同一个张量出现两次的二元运算，两边各取一次读锁，读锁可以同时持有，不会自己锁死自己。
和 RefCell 比：Tensor 仍然是 Send + Sync，可以在线程之间传。
和 UnsafeCell 比：并发写和读不会是未定义行为，最多是等锁。

锁的两个入口是 storage() 和 storage_mut()，都是 pub(crate)，取锁后直接 unwrap()。对外只公开一个只读的 storage_and_layout()。取写锁的公开方法不多：const_set / zero_set / one_set、slice_set、scatter_set、scatter_add_set、inplace_op1/2/3、Var::set。

「同一个张量用两次」在读读组合下没事，在写读组合下就是注释说的死锁。lock 这个实验写了一个 InplaceOp2（把第二个参数加到第一个上），先对两个不同的张量调，再让一个张量对自己调，后者放在子线程里，主线程等 3 秒（debug）：

let a = Tensor::arange(0f32, 4., &Device::Cpu)?;
println!("a + a = {:?}", (&a + &a)?.to_vec1::<f32>()?);
let b = Tensor::ones(4, DType::F32, &Device::Cpu)?;
a.inplace_op2(&b, &AddInto)?;
// 子线程里：a2.inplace_op2(&a2, &AddInto)

a + a = [0.0, 2.0, 4.0, 6.0]
a after inplace_op2(&b) = [1.0, 2.0, 3.0, 4.0]
inplace_op2(&self): no result after 3 s (Timeout)
slice_set(self view): Some("cannot use slice_set when self and src share their storage")

inplace_op2 先 self.storage_mut() 拿写锁，再 rhs.storage() 拿读锁；两个张量是同一把锁时，读锁永远等不到。slice_set、scatter_set、Var::set 在取锁之前先用 same_storage（比较两个 RwLock 的地址）挡掉这种情况，返回错误；inplace_op1/2/3 没有这道检查。自定义算子见自定义算子。

📝NOTE

冷知识：二十分钟里换了两次

最早的 Tensor_ 里是 storage: Arc<Storage>，数据不可变。2023-07-13 这天上午连着合并了三个 PR：

11:04，PR #154（「Tensor mutability」，提交说明「Use a ref-cell to provide tensor mutability」）换成 Arc<RefCell<Storage>>，取 storage 的方法返回 Result，借用冲突时报新增的 Error::BorrowError。上面那段注释也是这时加的，开头还写着「Storage uses a mutex here」，和代码对不上。
11:11，PR #155（「Fix a comment」）改掉开头那句，并补上 RefCell 那一条：「the resulting tensors would not be send or sync」。
11:25，PR #156（「Use a rwlock for inner mutability」）换成 Arc<RwLock<Storage>>，删掉 BorrowError，storage() 不再返回 Result。

#155 补上的那一条写的就是 RefCell 的缺点，14 分钟后 RefCell 被换掉（#156 的提交说明没有写理由）。三个提交都早于本地仓库里最早的 tag，git tag --contains 只能查到 0.5.1。

共享还是复制

share 这个实验对同一个 (3, 4) 的 f32 张量做一串操作，用 storage_and_layout() 打印每个结果底层数据的地址（CPU 上是 Vec 的指针，Metal 上是 Buffer 的地址）和布局：

== Cpu
             a: cpu 0x103397e00 len=12 shape=[3, 4] stride=[4, 1] off=0 id=TensorId(2)
         clone: cpu 0x103397e00 len=12 shape=[3, 4] stride=[4, 1] off=0 id=TensorId(2)
       reshape: cpu 0x103397e00 len=12 shape=[4, 3] stride=[3, 1] off=0 id=TensorId(3)
        narrow: cpu 0x103397e00 len=12 shape=[3, 2] stride=[4, 1] off=1 id=TensorId(4)
             t: cpu 0x103397e00 len=12 shape=[4, 3] stride=[1, 4] off=0 id=TensorId(5)
  t.contiguous: cpu 0x751b070840 len=12 shape=[4, 3] stride=[3, 1] off=0 id=TensorId(7)
     t.reshape: cpu 0x751b070870 len=12 shape=[12] stride=[1] off=0 id=TensorId(9)
  broadcast_as: cpu 0x103397e00 len=12 shape=[5, 4] stride=[0, 1] off=0 id=TensorId(11)
     unsqueeze: cpu 0x103397e00 len=12 shape=[1, 3, 4] stride=[4, 4, 1] off=0 id=TensorId(12)
        detach: cpu 0x103397e00 len=12 shape=[3, 4] stride=[4, 1] off=0 id=TensorId(2)
    contiguous: cpu 0x103397e00 len=12 shape=[3, 4] stride=[4, 1] off=0 id=TensorId(2)
 to_dtype same: cpu 0x103397e00 len=12 shape=[3, 4] stride=[4, 1] off=0 id=TensorId(2)
          copy: cpu 0x751b074f00 len=12 shape=[3, 4] stride=[4, 1] off=0 id=TensorId(13)
        affine: cpu 0x751b074270 len=12 shape=[3, 4] stride=[4, 1] off=0 id=TensorId(14)

三类结果：

类别	操作	新 `Tensor_`	新 storage
原样返回	`clone`；无事可做时的 `contiguous`、`detach`、`to_dtype`（同 dtype）、`to_device`（同设备）、`narrow`（取整维）、`transpose`（同一维）	否，`id` 不变	否
视图	`narrow`、`transpose` / `t`、`permute`、`broadcast_as`、`unsqueeze` / `squeeze`、连续张量的 `reshape`、`detach`	是，新 `id`	否，`self.storage.clone()`
复制或计算	所有算子（`affine`、`matmul`……）、非连续张量的 `contiguous` 和 `reshape`、`copy`、`to_dtype`、`to_device`	是	是

视图类方法都是同一个写法，以 transpose 为例：

        let op = BackpropOp::new1(self, |t| Op::Transpose(t, dim1, dim2));
        let tensor_ = Tensor_ {
            id: TensorId::new(),
            storage: self.storage.clone(),
            layout: self.layout.transpose(dim1, dim2)?,
            op,
            is_variable: false,
            dtype: self.dtype,
            device: self.device.clone(),
        };
        Ok(Tensor(Arc::new(tensor_)))

self.storage.clone() 只加内层 Arc 的计数，layout 换成新算的。detach 也属于这一类：它要一个 op 为 None 的新节点，所以必须有新的 Tensor_，但数据照样共享。reshape 分两支：输入连续时是视图（Layout::contiguous_with_offset(shape, self.layout.start_offset())），不连续时分配新 storage 再 copy_strided_src，所以上面 t.reshape 的地址变了。

视图共享 storage，写入就会穿透。share 里对一个 (2, 3) 张量的一列视图调 zero_set，再打印原张量，结果见文末动脑筋答案。

❗IMPORTANT

动脑筋：经视图清零

let a = Tensor::arange(0f32, 6., &device)?.reshape((2, 3))?;
a.narrow(1, 1, 1)?.zero_set()?;
println!("{:?}", a.to_vec2::<f32>()?);

a 打印出来是什么？CPU 和 Metal 一样吗？

copy 复制多少

copy 的文档说它和 clone 不同，「copies the actual storage」：

    pub fn copy(&self) -> Result<Tensor> {
        let op = BackpropOp::new1(self, Op::Copy);
        let tensor_ = Tensor_ {
            id: TensorId::new(),
            storage: Arc::new(RwLock::new(self.storage().try_clone(self.layout())?)),
            layout: self.layout.clone(),
            // ...

它把复制交给 Storage::try_clone，布局原样保留。两个后端的 try_clone 都忽略传进来的 layout：CPU 上是 CpuStorage::clone，整个 Vec 复制一份；Metal 上是 MetalStorage::try_clone：

    fn try_clone(&self, _: &Layout) -> Result<Self> {
        Ok(self.clone())
    }

MetalStorage 的 buffer 字段是 Arc<Buffer>，clone 只加计数。share 的两段输出把这两点都量出来了：

== Cpu
...
after copy().zero_set: a = [[0.0, 1.0, 2.0], [3.0, 4.0, 5.0]]
copy of narrow view: cpu 0x7920808000 len=1000000 shape=[2] stride=[1] off=10
== Metal(MetalDevice(DeviceId(1)))
             a: metal 0x751b039210 shape=[3, 4] stride=[4, 1] off=0 id=TensorId(25)
...
          copy: metal 0x751b039210 shape=[3, 4] stride=[4, 1] off=0 id=TensorId(36)
...
after copy().zero_set: a = [[0.0, 0.0, 0.0], [0.0, 0.0, 0.0]]
copy of narrow view: metal 0x751b039850 shape=[2] stride=[1] off=10

CPU 上 copy 是真复制，但复制的是整块 storage：对一个 100 万元素张量的 2 元素视图调 copy，得到的是 100 万元素的 Vec，偏移仍是 10。要一份紧凑的副本，用 force_contiguous（总是按布局拷成连续的新 storage）。
Metal 上 copy 的结果和原张量地址相同（都是 0x751b039210），对副本 zero_set 把原张量也清零了。也就是说，在 Metal 上 copy 和视图没有区别，文档承诺的「复制」不成立。candle-core、candle-nn、candle-transformers 的 src/ 里没有调用 Tensor::copy 的地方，所以模型推理不受影响；用户代码如果靠 copy 留一份快照再原地改，在 Metal 上会改到快照。

📝NOTE

没有蠢问题

问： storage 明明可以原地写，为什么 add、mul 这些算子每次都分配新的 storage，不在左操作数上原地算？

答：原地写会破坏计算图。BackpropOp::new2(self, rhs, ..) 把两个输入张量存进结果的 op 里，反向传播时还要读它们的值（Mul 的梯度要用到另一个输入）；如果 a.add(&b) 把结果写回 a 的 storage，图里存着的 a 就被改掉了。tensor.rs 里二元算子定义的上方有一句 TODO：「Also make an inplace version or a pre-allocated? This could be tricky if this can create cycles in the compute graph.」candle 把原地写限制在不参与求导的少数 API（const_set、slice_set、inplace_op*、Var::set）上。

clone 的代价

share 最后一段在 CPU 上量三件事（1024×1024 f32，release）：

clone: 4.0 ns
t(): 49.2 ns
t().contiguous() 1024x1024 f32: 2165.6 us

clone 是一次原子加，4 ns，和张量多大无关。t() 要新建一个 Tensor_：分配外层 Arc、为新的形状和步长各分配一个 Vec、从计数器取一个 id，约 50 ns，仍然和数据量无关。真正和数据量有关的是 contiguous() 这类复制：4 MB 的转置拷成连续要 2.2 ms，原因见布局与视图。

所以计算图把输入张量存成 Tensor（而不是引用）几乎没有成本，这正是 Tensor 上方那句注释说的「cloning is cheap when building the op graph」。

动脑筋答案

after narrow(1,1,1).zero_set: [[0.0, 0.0, 2.0], [3.0, 0.0, 5.0]]

CPU 和 Metal 都打印这一行。a.narrow(1, 1, 1) 是一个视图：形状 (2, 1)、步长 [3, 1]、起始偏移 1，和 a 共用 storage。const_set 取视图的写锁（也就是 a 的写锁），按视图的布局只写偏移 1 和 4 两个位置，所以 a 的第二列变成 0，其他元素不动。

小结

Tensor 是 Arc<Tensor_>，8 字节；clone 4 ns。Tensor_ 有 id、storage、layout、op、is_variable、dtype、device 七个字段；dtype 和 device 多存一份，dtype()、device() 不用取锁。
数据在 Arc<RwLock<Storage>> 里：读读可以并存，所以 a + a 没问题；写读同一把锁会死锁，slice_set 等方法用 same_storage 挡住，inplace_op* 没有挡。
视图（narrow、transpose、permute、broadcast_as、连续张量的 reshape、detach）新建 Tensor_、共享 storage，经视图写入会改到原张量；算子、contiguous（非连续时）、to_dtype、to_device 分配新 storage。
copy 在 CPU 上复制整块 storage（不管视图多小），在 Metal 上只复制 Arc<Buffer>，副本和原张量共用数据。