Rust in Everything
基于 burn 0.22.1

张量

这一部分讲 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。

  1. Tensor 与 Kind:Tensor<D, K> 只有一个字段,是一块 384 字节、按 64 对齐的不透明存储;Float / Int / Bool 是封闭的种类,Basic / Numeric / Ordered / FloatMath 这些公开 trait 决定一个方法对哪种张量可用;方法按值消耗 self,clone 只加一个引用计数,Flex 上唯一持有的左操作数可以原地写。
  2. BridgeTensor 与 *_impl:泛型方法体里只做检查和形状计算,真正的调用外提到不泛型的函数,它们只经手 BridgeTensor;obfuscate! 让 BridgeTensor 的定义不提 DispatchTensor。release 下的对照实验里,同一段用户代码换成直接持有 DispatchTensor,叶子 crate 的函数定义从 48 个涨到 450 个。
  3. Device:Device 是包着 DispatchDevice 的门面;构造函数随 feature 出现,Default 按固定顺序挑一个已开启的后端,一个也没有就 panic;相等只比硬件,不看 autodiff;to_device 保留张量自己的 autodiff 标记。
  4. 检查与报错: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。

💬 关联讨论

发起讨论
加载中...