Storage 与后端 trait

主线位置：算子主线的中段。Tensor::matmul 检查完形状，调用 self.storage().matmul(&rhs.storage(), (batching, m, n, k), ..)（读锁见 Tensor 的结构）。本页讲 Storage::matmul 这一层：先查两边是不是同一设备、同一 dtype，再按 (Cpu, Cpu) / (Metal, Metal) 匹配，交给 CpuStorage::matmul（CPU 后端）或 MetalStorage::matmul（Metal 后端）。

flowchart TB
  t["Tensor matmul"] --> s["Storage matmul"]
  s --> dev{"same_device"}
  dev -->|否| e1["DeviceMismatch"]
  dev --> dt{"same_dtype"}
  dt -->|否| e2["DTypeMismatch"]
  dt --> m{"match 设备对"}
  m -->|"Cpu Cpu"| cpu["CpuStorage matmul"]
  m -->|"Metal Metal"| mtl["MetalStorage matmul"]
  m -->|"Cuda Cuda"| cu["CudaStorage matmul"]
  cpu --> w["包回 Storage Cpu"]
  mtl --> w2["包回 Storage Metal"]

后端这一层有两样东西：一个枚举 Storage，在运行期决定走哪个后端；两个 trait BackendStorage / BackendDevice，在编译期保证每个后端都实现了同一组方法。trait 不用来做动态分派，分派全靠枚举上的 match。

本页的实验 crate（M2 Max、macOS 27、rustc 1.99.0，debug）：exp-candle/backend-ops（candle-core 0.11.0，features metal accelerate）和 exp-candle/backend-noaccel（不开任何 feature）。

一个枚举，三个变体

Storage 只有三个变体：

// We do not want to implement Clone on Storage as cloning may fail because of
// out of memory. Instead try_clone should be used.
#[derive(Debug)]
pub enum Storage {
    Cpu(CpuStorage),
    Cuda(CudaStorage),
    Metal(MetalStorage),
}

三个变体不受 feature 影响，永远都在。没开 metal 时，MetalStorage 这个名字指向另一个类型：lib.rs 按 cfg 二选一地 pub use：

#[cfg(feature = "metal")]
pub use metal_backend::{MetalDevice, MetalError, MetalStorage};

#[cfg(not(feature = "metal"))]
pub use dummy_metal_backend::{MetalDevice, MetalError, MetalStorage};

CUDA 同理（pub use dummy_cuda_backend as cuda;）。这样 storage.rs 里每个 match 都可以照写三支，不用到处加 #[cfg]。代价是没开的后端也占着一个变体，它的方法全是报错，见下文「没开 feature 时」。

Storage 没有实现 Clone。注释给的理由是复制可能因为内存不够而失败，所以用返回 Result 的 try_clone。Storage::try_clone 在 CPU 上是 Vec 的深拷贝；在 Metal 上交给 MetalStorage::try_clone，它的实现是 Ok(self.clone())，只复制 Arc<Buffer>，两边共用同一块显存。Tensor::copy 就建立在 try_clone 上，所以在 Metal 上「复制」出来的张量和原张量共享数据，Tensor 的结构一页有实测。

先检查，再按设备对匹配

Storage 上的方法（全是 pub(crate)，用户碰不到）都是同一个模子。以 Storage::matmul 为例：

    pub(crate) fn matmul(
        &self,
        rhs: &Self,
        bmnk: (usize, usize, usize, usize),
        lhs_layout: &Layout,
        rhs_layout: &Layout,
    ) -> Result<Self> {
        self.same_device(rhs, "matmul")?;
        self.same_dtype(rhs, "matmul")?;
        match (self, rhs) {
            (Self::Cpu(lhs), Self::Cpu(rhs)) => {
                let storage = lhs.matmul(rhs, bmnk, lhs_layout, rhs_layout)?;
                Ok(Self::Cpu(storage))
            }
            (Self::Cuda(lhs), Self::Cuda(rhs)) => {
                let storage = lhs.matmul(rhs, bmnk, lhs_layout, rhs_layout)?;
                Ok(Self::Cuda(storage))
            }
            (Self::Metal(lhs), Self::Metal(rhs)) => {
                let storage = lhs.matmul(rhs, bmnk, lhs_layout, rhs_layout)?;
                Ok(Self::Metal(storage))
            }
            (lhs, rhs) => Err(Error::DeviceMismatchBinaryOp {
                lhs: lhs.device().location(),
                rhs: rhs.device().location(),
                op: "matmul",
            }
            .bt()),
        }
    }

三步：

1. same_device 比较两边的 DeviceLocation；Metal 例外，要求是同一个 MetalDevice 实例（比较 DeviceId），同一块 GPU 上两次 Device::new_metal(0) 得到的设备也算不同。这条规则的来由见 DType 与 Device。
2. same_dtype 比较 DType，不做任何隐式转换。
3. match (self, rhs) 只列出同设备的组合，剩下的 (lhs, rhs) 再报一次设备不匹配。按 binary_impl 里的注释，这一支「Should not happen because of the same device check above but we're defensive anyway」。

结果统一包回同一个变体，所以 Storage::matmul 的返回值还是 Storage，Tensor::matmul 拿它和新形状、BackpropOp 拼成新张量（见计算图怎么记下来）。

只有一个输入的方法，比如 Storage::unary_impl，没有检查，直接按变体转发。

backend-ops 里的 storage 程序把这几种失败都触发一遍：

let a = Tensor::ones((2, 2), DType::F32, &cpu)?;
let b = Tensor::ones((2, 2), DType::F32, &metal)?;
show("cpu + metal", a.add(&b).map(|t| t.dims().to_vec()));
show("cpu matmul metal", a.matmul(&b).map(|t| t.dims().to_vec()));
let m2 = Device::new_metal(0)?;
let c = Tensor::ones((2, 2), DType::F32, &m2)?;
show("metal#1 + metal#2", b.add(&c).map(|t| t.dims().to_vec()));

cpu + metal: Err(device mismatch in add, lhs: Cpu, rhs: Metal { gpu_id: 4294968354 })
cpu matmul metal: Err(device mismatch in matmul, lhs: Cpu, rhs: Metal { gpu_id: 4294968354 })
metal#1 + metal#2: Err(device mismatch in add, lhs: Metal { gpu_id: 4294968354 }, rhs: Metal { gpu_id: 4294968354 })

第三行两边的 gpu_id 一模一样，报错却说不匹配：报错里打印的是 location()（GPU 的 registryID），判断用的却是 DeviceId。

两个 trait

每个后端要实现的方法写在 backend.rs 的两个 trait 里。BackendStorage 有 31 个方法：

BackendDevice 有 13 个方法：new、location、same_device、zeros_impl、alloc_uninit、storage_from_slice、storage_from_cpu_storage（和 _owned 版）、rand_uniform、rand_normal、set_seed、get_current_seed、synchronize，以及关联类型 Storage。

这里面没有 softmax、layer norm、注意力，也没有量化矩阵乘。前者在 candle-nn 里用自定义算子实现；后者走 quantized 模块自己的 QStorage（见 QTensor 与 QMatMul）。trait 的面越小，加一个后端要写的东西越少；代价是每个融合算子都得在 candle-nn 里为每个后端各写一遍。

两个 trait 没有默认实现，一个方法漏写就编译不过。但它们只是「契约」，代码里没有 dyn BackendStorage，也写不出来：

use candle_core::backend::BackendStorage;
use candle_core::cpu_backend::CpuDevice;

fn main() {
  let _v: Vec<Box<dyn BackendStorage<Device = CpuDevice>>> = Vec::new();
}

error[E0038]: the trait `BackendStorage` is not dyn compatible
 --> examples/dyn_storage.rs:5:23
  |
5 |   let _v: Vec<Box<dyn BackendStorage<Device = CpuDevice>>> = Vec::new();
  |                       ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^ `BackendStorage` is not dyn compatible
  |
  = note: the trait is not dyn compatible because it requires `Self: Sized`
  = note: for a trait to be dyn compatible it needs to allow building a vtable
          for more information, visit <https://doc.rust-lang.org/reference/items/traits.html#dyn-compatibility>
  = help: the following types implement `BackendStorage`:
            candle_core::CpuStorage
            candle_core::CudaStorage
            candle_core::dummy_metal_backend::MetalStorage
            candle_core::MetalStorage
          consider defining an enum where each variant holds one of these types,
          implementing `BackendStorage` for this new enum and using it instead

BackendStorage: Sized 一条就排除了 trait 对象；就算去掉它，unary_impl<B> 这类带类型参数的方法也进不了虚表。rustc 最后给的建议「consider defining an enum」正是 candle 的做法。列表里同时出现了两个 MetalStorage：dummy_metal_backend 模块在 lib.rs 里无条件编译，开了 metal 也在，只是名字没有导出到 crate 根。

每个算子一个零大小类型

binary_impl<B: BinaryOpT> 的 B 不是值，是类型。op.rs 给每个逐元素算子定义一个空结构体（Add、Mul、Exp、Silu……），再为它实现 UnaryOpT 或 BinaryOpT：

pub trait UnaryOpT {
    const NAME: &'static str;
    const KERNEL: &'static str;
    const V: Self;
    fn bf16(v1: bf16) -> bf16;
    fn f16(v1: f16) -> f16;
    fn f32(v1: f32) -> f32;
    fn f64(v1: f64) -> f64;
    fn u8(v1: u8) -> u8;
    // ...
    fn f8e4m3(v1: f8e4m3) -> f8e4m3;

    #[inline(always)]
    fn bf16_vec(xs: &[bf16], ys: &mut [bf16]) {
        xs.iter().zip(ys).for_each(|(&x, y)| *y = Self::bf16(x))
    }
    // ...
    #[inline(always)]
    fn f32_vec(xs: &[f32], ys: &mut [f32]) {
        xs.iter().zip(ys).for_each(|(&x, y)| *y = Self::f32(x))
    }
    // ...
}

trait 里有三类成员，分给三个用户：

• 每个 dtype 一个标量函数（f32(v1: f32) -> f32 等），CPU 后端在非连续的路径上逐元素调用。
• 每个浮点 dtype 一个切片函数（f32_vec），默认实现就是循环调标量函数，CPU 后端在连续块上调用；某个算子有更快的实现时覆盖它。
• 两个字符串常量：NAME 用在报错里，KERNEL 给 GPU 后端拼内核名（Exp 的 KERNEL 是 "uexp"，Metal 后端再按 dtype 映射到 unary.metal 里的 exp_f32 之类）。

实现由两个宏批量生成。unary_op! 的第二种写法多带两个函数名，开了 accelerate（或 mkl）时用它们覆盖 f32_vec / f64_vec：

            #[cfg(feature = "accelerate")]
            #[inline(always)]
            fn f32_vec(xs: &[f32], ys: &mut [f32]) {
                crate::accelerate::$f32_vec(xs, ys)
            }

unary_op!(Exp, "exp", v, v.exp(), vs_exp, vd_exp);
unary_op!(Log, "log", v, v.ln(), vs_ln, vd_ln);
// ...
unary_op!(Neg, "neg", v, -v);

所以 Exp 在 accelerate 构建里的 f32_vec 是一次 vvexpf 调用，Neg 永远是 Rust 循环。二元算子的 bin_op! 同理，目前只有 Add 带了覆盖（bin_op!(Add, add, |v1, v2| v1 + v2, vec_add);），指向 cpu::kernels::VecOps::vec_add。

用户调用 a.exp() 时，tensor.rs 的 unary_op! 生成的方法写的是 self.storage().unary_impl::<crate::op::Exp>(self.layout())，类型在这里就定死了。之后 Storage::unary_impl::<Exp> → CpuStorage::unary_impl::<Exp> → unary_map_vec(.., Exp::f32, Exp::f32_vec) 全是单态化（monomorphization）出来的专用代码，没有函数指针、没有按算子的 match。代价是二进制体积：算子数 × 后端数 × dtype 数个函数体。同一个算子在计算图里还要再记一次，那里用的是运行期的值 UnaryOp::Exp（op.rs 里的 UnaryOp 枚举），因为反向传播要在运行期按节点 match。

整数 dtype 是这个设计的一个缝。unary_op! 给 u8、u32、i16、i32、i64 生成的是 todo!("no unary function for u32")，而 CpuStorage::unary_impl 遇到整数存储照样调用它，前面没有任何 dtype 检查。backend-noaccel 里的 dummy 程序：

let t = Tensor::new(&[1u32, 2, 3], &Device::Cpu).unwrap();
let r = t.exp();

thread 'main' (3709843) panicked at ~/.cargo/registry/src/index.crates.io-1949cf8c6b5b557f/candle-core-0.11.0/src/op.rs:584:1:
not yet implemented: no unary function for u32

在一个几乎每个方法都返回 Result 的库里，这是 panic 而不是 Err，位置指向 unary_op!(Exp, ..) 那一行宏调用。

没开 feature 时：dummy 后端

dummy_metal_backend.rs 里 MetalDevice 和 MetalStorage 都是没有字段的单元结构体，两个 trait 照样实现，方法体分两种：

macro_rules! fail {
    () => {
        unimplemented!("metal support has not been enabled, add `metal` feature to enable.")
    };
}

• 返回 Result 的方法（matmul、to_cpu_storage、BackendDevice::new 等）一律 Err(Error::NotCompiledWithMetalSupport)。
• 不返回 Result 的方法（dtype、device、location、same_device）只能调 fail!()，也就是 panic。

正常的入口是 Device::new_metal，它返回 Err，用户拿不到一个 dummy 的 Metal 张量。不过单元结构体可以手写出来：Device::Metal(candle_core::MetalDevice) 在不开 metal 时也能编译，再调 .location() 就落进 panic 分支，报「not implemented: metal support has not been enabled, add metal feature to enable.」。正常用法碰不到这些分支，但它们不是不可达的。synchronize 是唯一返回 Ok(()) 的方法。backend-noaccel（不开 feature）里：

metal_is_available=false
new_metal: Err(the candle crate has not been built with metal support)

和 burn 的 Dispatch 对照：两边都是「枚举 + 每个后端实现同一个 trait」，差别在枚举放在哪一层。burn 0.22 的 DispatchTensor 是用户可见的张量类型，变体由 build.rs 算出的 cfg 裁剪掉，一个后端都不开时只剩不可构造的 NoBackend，分派代码由过程宏 #[backend_dispatch] 生成；candle 的 Storage 藏在 Tensor 里面，变体永远是三个，没开的后端用 dummy 类型占位，storage.rs 里 849 行 match 是手写的。burn 的 trait 有几百个方法、一部分带默认实现（「用别的 op 拼出来」）；candle 的两个 trait 合计 44 个方法，没有默认实现，融合算子放到 trait 外面。

动脑筋答案

# backend-noaccel（不开 feature）
size_of MetalStorage = 0
size_of Storage = 32
size_of CpuStorage = 32

# backend-ops（开 metal）
size_of MetalStorage = 96
size_of MetalDevice = 72
size_of CudaStorage = 0
size_of CpuStorage = 32
size_of Storage = 96

• 不开 feature：CpuStorage 是 24 字节的 Vec 加 8 字节对齐后的判别值，共 32；两个 dummy 类型是 0 字节，Storage 的判别值塞进 CpuStorage 判别值没用到的取值里（niche），整体还是 32。
• 开 metal：MetalStorage 有 96 字节，其中 MetalDevice 占 72（一个 DeviceId、一个 Device 句柄和七个 Arc），Storage 跟着变成 96。这对每个张量是一次性的，Storage 放在 Arc<RwLock<..>> 里，Tensor 本身大小不变。
• 没开的后端在类型里不占空间，这是用 dummy 类型而不是删掉变体也不吃亏的原因。

小结

• Storage 永远是 Cpu / Cuda / Metal 三个变体；没开的后端由 lib.rs 按 cfg 换成 dummy 类型，match 不用写 #[cfg]。
• 两个输入的方法先 same_device、same_dtype，再按设备对 match；Metal 判断「同设备」看实例，报错里打印的却是 GPU 的 registryID。
• BackendStorage（31 个方法）和 BackendDevice（13 个）是编译期契约，要求 Sized，写不成 trait 对象；分派全靠枚举。
• 每个逐元素算子是一个零大小类型，UnaryOpT / BinaryOpT 提供每个 dtype 的标量函数、切片函数和内核名；开 accelerate 时宏把切片函数换成 Accelerate 调用。整数 dtype 的一元运算是 todo!()，u32 张量调 exp() 直接 panic。
• dummy 后端的 Result 方法返回 NotCompiledWithMetalSupport，用户能碰到的只有 Device::new_metal 的那个 Err。