张量
这一部分讲 candle-core 里用户直接打交道的那一层:a.matmul(&b)? 的 a 是什么,它的数据怎么摆,它在哪个设备、是什么元素类型,出错时返回什么。调用越过这一层、进入 Storage 按设备分派之后的事,在后端。
代码集中在 candle-core/src/ 下的几个文件:tensor.rs(3116 行,Tensor 的几乎所有方法)、layout.rs 和 shape.rs(形状、步长、偏移)、strided_index.rs(遍历非连续数据的迭代器)、dtype.rs、device.rs、error.rs。
- Tensor 的结构:
Tensor是Arc<Tensor_>,数据在第二层Arc<RwLock<Storage>>里;clone只加引用计数(实测 4 ns);narrow、transpose、broadcast_as和连续张量的reshape共享 storage,经过视图写入会改到原张量;Metal 上的copy()其实不复制。 - 布局与视图:
Layout是形状、按元素计的步长、起始偏移;视图操作只算新的Layout;后端用StridedIndex、StridedBlocks(0.11.0 新增UniformBlocks)遍历非连续数据;同一个转置张量,matmul 几乎不受影响,逐元素运算慢 1.5 到 15 倍;Metal 的批量 matmul 遇到某些步长会算错。 - DType 与 Device:
DType有 14 个变体,其中 4 个只能存不能算;WithDType把 Rust 类型接到DType上;Device是三个变体的枚举,Metal 按设备实例而不是物理 GPU 判断「同一设备」;不同 dtype、不同设备之间不做隐式转换,直接报错;Metal 上真正能算的浮点类型只有 F32、F16、BF16。 - 错误:一个
Error枚举、一个Result别名;形状、dtype、设备不匹配各有固定的报错格式;bail!构造Msg;backtrace 由Backtrace::capture按环境变量决定,打开后每个错误多花约 11 µs,还会改变match的结果;from_vec从 0.9.0 起不再检查元素个数。
flowchart TB user["a.matmul b"] --> chk["形状检查"] chk -->|失败| err["Err ShapeMismatch"] chk --> lock["取 storage 读锁"] lock --> st["Storage matmul"] st --> dev["同设备 同 dtype"] dev -->|失败| err2["Err Mismatch"] dev --> be["后端按 Layout 计算"] be --> out["新 Tensor"]
第一页讲张量本身的结构,第二页讲数据在存储里怎么摆、后端怎么按步长读,第三页讲后端分派前的两道检查(设备、dtype),第四页讲这一路上出错时用户拿到什么。Tensor_ 里的 op 字段怎么记下计算图,见计算图怎么记下来。
💬 关联讨论
发起讨论加载中...