张量
这一部分讲 burn-tensor:用户写下 a.matmul(b) 之后、调用落进 Dispatch 之前的那一段。0.22 把后端从 Tensor 的类型参数里拿掉了(Tensor<B, D, K> 变成 Tensor<D, K>),后端改由张量所在的 Device 在运行期决定。于是 burn-tensor 只剩三件事要做:在类型上区分秩和种类,在入口处检查形状和设备,把调用换成不带泛型的形式交给 burn-dispatch。
代码集中在 crates/burn-tensor/src/ 下:tensor/api/ 是公开的方法(base.rs 一个文件近 4000 行),bridge/ 是连接 Dispatch 的内部 trait 和 BridgeTensor,device.rs 是 Device 门面,tensor/api/check.rs 是 TensorCheck。TensorData、DType、Shape 这些和后端无关的类型来自 burn-std。
- Tensor 与 Kind:
Tensor<D, K>只有一个字段,是一块 384 字节、按 64 对齐的不透明存储;Float/Int/Bool是封闭的种类,Basic/Numeric/Ordered/FloatMath这些公开 trait 决定一个方法对哪种张量可用;方法按值消耗self,clone只加一个引用计数,Flex 上唯一持有的左操作数可以原地写。 - BridgeTensor 与 *_impl:泛型方法体里只做检查和形状计算,真正的调用外提到不泛型的函数,它们只经手
BridgeTensor;obfuscate!让BridgeTensor的定义不提DispatchTensor。release 下的对照实验里,同一段用户代码换成直接持有DispatchTensor,叶子 crate 的函数定义从 48 个涨到 450 个。 - Device:
Device是包着DispatchDevice的门面;构造函数随 feature 出现,Default按固定顺序挑一个已开启的后端,一个也没有就 panic;相等只比硬件,不看 autodiff;to_device保留张量自己的 autodiff 标记。 - 检查与报错:
TensorCheck把一个 op 的所有问题攒在一起再 panic,release 下也开着;cat不查设备,切片越界会被截断而不报错;读数据的try_*方法返回ExecutionError/TensorReadError,其余都是 panic。
flowchart TB user["a.matmul b"] --> chk["TensorCheck 检查"] chk -->|失败| boom["panic 原文"] chk -->|通过| kind["K matmul 不泛型"] kind --> bridge["BridgeTensor into_parts"] bridge --> disp["Dispatch float_matmul"] dev["Device"] -.->|决定后端| disp
前两页讲「张量是什么、调用怎么离开泛型」,第三页讲「交给哪个后端由什么决定」,第四页讲「哪些错误在这一层就被拦下」。Dispatch 拿到调用之后怎么按设备 match 到具体后端,见 Dispatch;后端要实现哪些 op,见 Backend trait。
💬 关联讨论
发起讨论加载中...