Rust in Everything
基于 tenferro-runtime 0.8.0

数据模型

主线位置:直接计算主线的第一步,TypedTensor::<f64>::from_vec_col_major(vec![2, 2], data)?。本页讲这一行造出来的是什么:f64 怎么变成运行期的 DType::F64,类型里看不见的另外两个参数 R 和 D 是什么,它和动态的 Tensor 怎么互转,哪些错误在编译期就挡住、哪些要等到运行期。数据怎么排在下一页 列主序与视图,内存归谁在 存储与所有权。

flowchart TB
  table["预置标量表"] --> dtype["DType 标签"]
  table --> spec["MemberSpec 提升"]
  scalar["TensorScalar"] --> typed["TypedTensor T R D"]
  rank["R 秩"] --> typed
  repr["D 表示"] --> typed
  typed -->|Host| hostd["Vec 加回收器"]
  typed -->|Gpu| grp["AllocationGroup"]
  typed -->|Dynamic| either["二者之一"]
  dynt["Tensor"] --> preset["七个变体"]
  preset --> typed
  dynt -->|External| ext["ErasedHostTensor"]

本页的实验 crate 是 exp-tenferro8/tensor-exp(tenferro-runtime、tenferro-tensor、tenferro-tensor-core、tenferro-cpu 等都是 0.8.0,默认 native provider),程序 data_model、eager_promote、custom_scalar、alloc_count;另有一个只用来看编译错误的 tensor-cfail。全部 release,M2 Max、macOS 27、rustc 1.99.0。

两个 crate 各管什么

tenferro-tensor-core(src/ 2040 行,不含测试)只剩与存储无关的元数据:dtype 标签、标量 trait、秩、TensorLayout 和布局校验。它的 crate 文档写明「It does not own tensor storage」。tenferro-tensor(src/ 26273 行,不含 tests 目录和 *tests.rs,其中 types.rs 一个文件 11680 行)放所有带存储的张量:TypedTensor、Tensor、各种视图、DefaultScalars、ErasedHostTensor。

0.7.1 时 core 里还有一个纯主机的 HostTensor 和一个叫 Tensor 的别名,和 tenferro-tensor 的张量是两套。0.8.0 把它们删了(见下面的冷知识),工作区里现在只有一个 Tensor。tenferro_tensor::core 模块的 转出列表 只剩 DType、DynRank、Rank、TensorLayout、ShapeVec 这类元数据;张量类型都从 tenferro_tensor 根(或 tenferro_runtime 的 prelude)导入。

名字仍有一处重复:两个 TensorScalar。core 的 TensorScalar 只有 Real 关联类型和 dtype();tenferro-tensor 的 TensorScalar 另带一组在 Tensor 和 TypedTensor 之间转换的钩子。两者都用私有的 Sealed 封住,实现者是同样七个类型:f32、f64、i32、i64、bool、Complex32、Complex64。用户代码平时只碰后一个。

dtype:一张表展开

DType 不是手写的枚举,是 define_scalar_tag! 宏的一次调用(引用时删去了文档注释):

define_scalar_tag! {
    pub enum DType {
        F32 => f32 : Float 0 32,
        F64 => f64 : Float 1 64,
        I32 => i32 : Integer 0 32,
        I64 => i64 : Integer 1 64,
        Bool => bool : Boolean 0 0,
        C32 => Complex32 : Complex 0 32,
        C64 => Complex64 : Complex 1 64,
    }
    external External(core::any::TypeId);
}

每行是「标签 => Rust 类型 : 种类 级别 位宽」。define_scalar_tag! 从这张表生成 DType 枚举(七个预置变体加一个 External(TypeId),TypeId 16 字节,所以实测 size_of::<DType>() 是 24)、每个标签的 spec(),以及按同一顺序排好的 TAGS 和 SPECS 两张常量表。

0.7.1 时同一个宏还顺带生成了装数据的载荷枚举。0.8.0 把这一步拆开:core 只生成标签,数据载荷归 tenferro-tensor。预置标量集 DefaultScalars 是手写的,七个变体各装一个 TypedTensor<_, DynRank, Host>;它的 ScalarSet 实现 最终落到 core 的 promote_in_set。给下游自定义标量集用的 define_scalar_set! 在 tenferro-tensor 里,内部再调 define_scalar_tag!;仓库里调用它的只有 ext/ 下 bf16、df64 等几个验证 crate。

提升规则只看 MemberSpec。promote_specs 先按种类配对:Bool 让位给另一边;整数遇上浮点,目标是「级别和位宽都是 u32::MAX 的浮点」;同种类取级别高的;浮点遇上复数取两者中较宽的位宽。promote_in_set 在表里找不到正好相等的 spec 时,挑同种类里位宽够用的最窄一个,都不够就挑最宽的。实测:

promote(I32, F32) = F64
promote(I64, F32) = F64
promote(F32, C32) = C32
promote(F64, C32) = C64
promote(Bool, I32) = I32

I32 和 F32 提升成 F64:目标位宽是 u32::MAX,没有哪种浮点「够宽」,只好取最宽的 F64。

📝NOTE
没有蠢问题

问: 既然有提升规则,为什么 f64 的 Tensor 和 f32 的 Tensor 做 matmul 会报错?

答: 规则只被少数调用方用。公开入口是 tenferro-tensor 的 promote_dtype,它转给 DefaultScalars::promote;tenferro-runtime 的 shape_infer::promote_dtype 只是转调它,eager 自动微分(tenferro-ad 的 eager_exec.rs)经 runtime 这一层在执行前提升输入。直接计算主线上的运算不做隐式转换。实测两个 Tensor(F64 和 F32)在 CpuEngine::serial() 上:x.matmul(&y, &mut engine) 返回 dot_general: dtype mismatch: expected F64, actual F32,x.add(&y, &mut engine) 返回 add: dtype mismatch: expected F64, actual F32;换成两个 EagerTensor,s.matmul(&x, &y) 和 s.add(&x, &y) 都成功,结果是 F64。eager 怎么记录这次提升,见 eager AD。

core 里还有两个和 dtype 并列的 trait。Scalar 的关联常量 DOMAIN 取 ScalarDomain 的 Field 或 NonField,说的是「这个标量的普通算术是不是域,求导规则能不能照搬」。七个预置类型里只有 bool 是 NonField。配套的 ad_admission 在 0.8.0 里仍没有返回 Ok 的分支:实测 ad_admission::<f64>(1) 是 Err(AdRuleUnavailable),ad_admission::<bool>(1) 是 Err(NonFieldScalar);工作区 crates/ 里除了 core 自己的转出没有调用方。f64 照样能求导(见 primitive AD);它是给下游自定义标量预留的接口。

TypedTensor 的三个参数

pub struct TypedTensor<T, R: TensorRank = DynRank, D: Representation = Dynamic> {
    shape: R::Shape,
    placement: Placement,
    /// Column-major element strides, cached at construction so a strided view of
    /// this tensor borrows its layout instead of recomputing it per operation.
    layout: StrideVec,
    storage: D::Storage<T, R>,
}

TypedTensor 的四个字段:形状、放置(内存种类、设备、CPU 域)、构造时算好的列主序步长,和由 D 决定的存储。类型本身对 T 没有任何约束:T: Clone 只在复制的构造函数上要,TensorScalar 只在真要用 dtype 或数值内核的方法上要。

R:秩

R 由 sealed trait TensorRank 约束,只有两种实现:

RR::ShapeR::Strides
DynRank(默认)ShapeVec = TinyVec<[usize; 8]>StrideVec = TinyVec<[isize; 8]>
Rank<N>[usize; N][isize; N]

ShapeVec 和 StrideVec 从 0.7.1 的 SmallVec 换成了 tinyvec 的 TinyVec,内联容量是常量 INLINE_RANK = 8,注释说它和 strided-view 的同名常量一致、由 tenferro-cpu-basic 里的编译期断言保证。秩不超过 8 时形状和步长不上堆。用计数分配器实测(alloc_count):from_vec_col_major 拿一个现成的形状 Vec 和数据 Vec,0 次分配;秩 10 的张量构造要 2 次(形状和步长各溢出一次)。

构造函数的形状参数是 impl IntoRankShape<R>。给 Rank<N> 实现这个 trait 的类型分两类(定长数组的实现在 rank.rs):[usize; N] 和 &[usize; N] 只在长度等于 N 时有实现,长度不对是编译错误;Vec<usize>、&[usize]、ShapeVec 对任意 N 都有实现,长度在运行期比。两种都实测过:

let a = TypedTensor::<f64, Rank<2>>::from_vec_col_major([2, 3, 1], vec![0.0; 6]).unwrap();
error[E0277]: the trait bound `[{integer}; 3]: IntoRankShape<Rank<2>>` is not satisfied
    --> src/bin/rank_array.rs:4:59
     |
   4 |   let a = TypedTensor::<f64, Rank<2>>::from_vec_col_major([2, 3, 1], vec![0.0; 6]).unwrap();
     |           ----------------------------------------------- ^^^^^^^^^ the trait `IntoRankShape<Rank<2>>` is not implemented for `[{integer}; 3]`
...
Rank<2> from vec![2,3,1] -> Err: from_vec_col_major: rank mismatch: expected 2, actual 3
try_into_rank::<3>() -> Err: tensor reinterpretation failed: TypedTensor::try_into_rank: rank mismatch: expected 3, actual 2
try_into_rank::<2>() same buffer = true
shape [2,3] with 5 values -> Err: from_vec_col_major: shape product 6 does not match data length 5

try_into_rank 把动态秩换成静态秩,缓冲区地址不变;失败时错误里带回原张量(「tensor reinterpretation failed」前缀来自 ReinterpretError)。元素个数和数据长度对不上永远是运行期错误。

静态秩的好处在下标上。0.8.0 起 TypedTensor<T, Rank<N>> 实现了 Index<[usize; N]>,host_col_major_view() 返回的 ColMajorView 也一样,越界就 panic;要可失败的版本用 get。实测 r2[[1, 2]] = 6、ColMajorView[[0, 2]] = 3(数据是 [[1,2,3],[4,5,6]] 的列主序)。上游的 trybuild 测试 static_rank_mismatch 固定了 Rank<2> 视图传给要 Rank<3> 的函数是 E0308。

D:表示

D 是 0.8.0 新加的参数,由 sealed trait Representation 约束,三个标记类型各自选一种存储:

DD::Storage含义主机访问
HostHostStorage<T>:一个 Vec<T>,外加可选的池回收器纯主机,静态保证as_slice() 不会失败;有 Clone、Index
GpuGroupStorage<R>:一个装箱的分配组交给分配组管的内存host_data() 返回 Result
Dynamic(默认)DynamicStorage:Host 或 Group 二选一运行期才知道host_data() 返回 Result

Gpu 的名字容易误读。它的文档注释写「The marker is not proof of a particular device ordinal or provider」:它说的是「这块内存由分配组持有」,不是「在 GPU 上」。实测把一个 Host 张量 promote() 成 Gpu,host_data() 照样成功、指针不变,is_backend_buffer() 是 false。

三种表示之间的转换都不动数据:

Host from_host_vec_col_major same Vec: true
Host clone new buffer: true, equal: true
Host index [&[1, 0]] = 2
into_dynamic same buffer: true
into_host same buffer: true
promote -> Gpu: host_data same buffer: true, is_backend_buffer: false
group-backed into_host -> Err: TypedTensor::into_host: runtime state failure: the tensor is group-backed rather than plain host storage; owner returned, shape [2, 2]

into_dynamic 只换标记;into_host / into_gpu 是带检查的收窄,失败时把原张量放在 ReinterpretError 里还回来。设计文档 tensor-session-redesign-1938.md 的 D5 把这条写成规矩:「failure must return it with the typed cause rather than silently destroy the input」。

默认的 from_vec_col_major 造的是 Dynamic 张量,里面放的是 Host 那一支:数据 Vec 直接收下,不建分配组。分配组只在需要时才出现(张量进了 AllocationGroup,或来自后端缓冲区),见 存储与所有权。

Host 不要求 T 是 tenferro 的标量。实测(custom_scalar)一个自定义的 #[derive(Clone)] struct Interval { lo: f64, hi: f64 } 可以放进 TypedTensor<Interval, Rank<2>, Host>,按 [usize; 2] 下标、借 ColMajorView、clone,再 into_dynamic;只有要 dtype 或数值内核的方法才要 TensorScalar。

📝NOTE
冷知识:HostTensor 是怎么没的

0.7.1 有两套拥有型张量:core 里纯主机的 HostTensor<T>(还有一个 Tensor 别名),和 tenferro-tensor 的 TypedTensor。issue #1938 的设计记录在「Critique incorporated」表里点了这件事:「HostTensor was both a core type and an alias to a higher-crate type」,结论是「One public tensor family in tenferro-tensor; core keeps metadata/scalar facilities, with no reverse dependency (D10)」。PR #1944 先给 TypedTensor 加了 Host/Gpu/Dynamic 参数,PR #1947 里的提交 4e15272a 再删掉旧类型,提交说明:「HostTensor / HostTensorView were a second owner/view family beside TypedTensor<T, R, Host>. Delete them with no alias」,并写明「tenferro_tensor::core re-exports metadata only」。两个 PR 都首次进入 v0.8.0。迁移就是把 HostTensor<T> 换成 TypedTensor<T, DynRank, Host>,HostTensorView 换成 TypedTensorView<'a, T, DynRank, Host>。

动态 Tensor:七个 TypedTensor 的枚举

Tensor 只有一个字段,类型是私有枚举 TensorPayload:

#[derive(Debug)]
enum TensorPayload {
    Native(PresetTensor),
    External(crate::ErasedHostTensor, Placement),
}

#[derive(Debug)]
enum PresetTensor {
    F32(TypedTensor<f32>),
    F64(TypedTensor<f64>),
    I32(TypedTensor<i32>),
    I64(TypedTensor<i64>),
    Bool(TypedTensor<bool>),
    C32(TypedTensor<Complex32>),
    C64(TypedTensor<Complex64>),
}

七种预置 dtype 走 Native,每个变体就是一个默认参数的 TypedTensor<T>;dtype 在枚举标签上,Tensor::dtype() 对标签做一次 match。External 给下游自定义标量:ErasedHostTensor 装任意 TypedTensor<U, DynRank, Host>,对应 DType::External(TypeId)。

0.7.1 的 as_typed 靠 #[repr(transparent)] 加一次 unsafe 指针转换。0.8.0 的 TypedTensor 不再是 repr(transparent),Tensor::as_typed 改成安全的 Any 向下转型:

    pub fn as_typed<T: TensorScalar>(&self) -> Option<&TypedTensor<T>> {
        if self.dtype() != T::dtype() {
            return None;
        }
        match &self.payload {
            TensorPayload::Native(preset) => with_preset!(preset, |typed| {
                (typed as &dyn Any).downcast_ref::<TypedTensor<T>>()
            }),
            TensorPayload::External(..) => None,
        }
    }

拥有型的两个方向只是搬字段:Tensor::from_typed 把张量包进对应变体,into_typed 匹配变体再拆出来。实测:

dynamic.dtype() = F64, shape = [2, 2]
as_typed::<f32>() is_none = true
as_typed::<f64>() same buffer = true
into_typed::<f32>() -> Err: TensorScalar::into_typed: dtype mismatch: expected F32, actual F64; owner returned dtype F64
into_typed::<f64>() same buffer = true

into_typed 失败时不再吞掉输入:返回的是 ReinterpretError<Tensor>,into_parts() 拿回原张量和错误,同一个值可以接着按别的 dtype 试。as_typed 只对 TypedTensor<T, DynRank, Dynamic> 有:要静态秩,再调 try_into_rank。

为什么要两种类型?TypedTensor 让 dtype 进类型系统,写数值代码时 f64 和 f32 混用是编译错误;Tensor 给图、运行时、扩展算子这些只能在运行期知道 dtype 的地方用,traced 程序的输入输出、AllocationGroup 的成员都是 Tensor。两者之间切换不碰数据。

大小与 Clone

TypedTensor<f64>                    280
TypedTensor<f64, DynRank, Host>     280
TypedTensor<f64, DynRank, Gpu>      248
TypedTensor<f64, Rank<2>>           216
TypedTensor<f64, Rank<2>, Host>     216
Tensor                              312
Option<Tensor>                      312
TypedTensorView<f64>                680
TypedTensorView<f64, DynRank, Host> 680
TypedTensorView<f64, Rank<2>>       552
TensorLayout<DynRank>               168
TensorLayout<Rank<2>>               40
ShapeVec                            80
Placement                           80
DType                               24
ErasedHostTensor                    232
  • TypedTensor<f64> 是 280 字节:形状 80、放置 80、缓存的步长 80,剩下 40 是存储(Vec 加可选的 Weak 回收器)。Gpu 的存储只是一个 Box,所以小 32 字节。0.7.1 时每个张量内联了一整个分配组,是 776 字节。
  • 步长缓存是 #2149 加的,提交 4a528b3a 的说明:「An owned tensor recomputed its column-major strides on every view ... two allocations per operation on the eager paths」,代价写得很直白:「+72 bytes per tensor for the cache」。实测 as_view、transpose_view、slice_view、reshape_view 都是 0 次分配。
  • Tensor 的 312 由 External 变体决定(ErasedHostTensor 232 加 Placement 80),不是由 TypedTensor 的 280 决定;标签藏进空位(niche),Option<Tensor> 也是 312。上游测试 erased_tensor_size_stays_within_the_documented_bound 的上限仍是 0.7.1 的 776,现在离它很远。
  • 视图比拥有型张量还大:TypedTensorView 有一个 root: Option<GroupReadView> 字段,里面是一整份描述符记录,Host 标记的视图也留着这个槽。设计文档承认这一点:D3 的「never carries an optional full group descriptor」「holds for a host view at runtime but not as a type-level absence of the slot」,去掉它要牺牲视图对生命周期的协变,所以记成未完成项。
❗IMPORTANT
动脑筋:TypedTensor<f64> 是 280 字节,TypedTensor<f64, Rank<2>> 是 216。TensorLayout<DynRank> 和 TensorLayout<Rank<2>> 差了 128 字节,为什么张量本身只差 64?

看结构体的四个字段,哪几个的类型随 R 变。

Clone 只给 Host。impl Clone for TypedTensor<T, R, Host> 就是 duplicate():深拷贝成一个新的普通 Vec。默认的 Dynamic 张量没有 Clone,上游的 trybuild 测试 typed_tensor_not_clone 固定了这条错误,实测一样:

error[E0599]: no method named `clone` found for struct `TypedTensor<T, R, D>` in the current scope
 --> src/bin/no_clone.rs:5:13
  |
5 |   let b = a.clone();
  |             ^^^^^ method not found in `TypedTensor<f64>`

理由在设计文档 D4:「Default Dynamic tensors do not implement Clone or infallible indexing that could silently download/map a device value」。Dynamic 的内容可能不在主机上,clone 或 t[&[i]] 一旦成立就得偷偷下载。Rank<N> 的 Dynamic 张量有 Index<[usize; N]>,它走 get,存储读不到时 panic;动态秩的 t[&[1][..]] 则是 E0608(cannot index into a value of type TypedTensor<f64>)。要第二份数据,调 duplicate();要共享,借一个视图;确定是主机数据,into_host() 之后就能 clone。

编译期与运行期

检查TypedTensor<T, Rank<N>>TypedTensor<T>Tensor
两个操作数 dtype 相同编译期(E0308)编译期(E0308)运行期,dtype mismatch
秩数组形状编译期;Vec 形状运行期运行期运行期
元素个数等于数据长度运行期运行期运行期
matmul 的内积维相等运行期运行期运行期
clone只有 Host 有;Dynamic 是 E0599同左没有
数据能不能直接在主机上读Host 编译期;其余运行期同左运行期

dtype 混用的编译错误长这样(b 是 TypedTensor<f32>):

error[E0308]: mismatched types
   --> src/bin/dtype_mix.rs:8:20
    |
  8 |   let c = a.matmul(&b, &mut engine).unwrap();
    |             ------ ^^ expected `&TypedTensor<f64>`, found `&TypedTensor<f32>`
    |             |
    |             arguments to this method are incorrect
...

形状一律是运行期的值,即使用了 Rank<N>,[2, 3] 和 [3, 2] 也是同一个类型。和 burn 对照:burn 0.22 的 Tensor<D, K> 把秩 D 写在类型上,dtype 留在运行期(见 Tensor 与 Kind);candle 的 Tensor 秩和 dtype 都在运行期(见 Tensor 的结构)。tenferro 两种都给,用 R 选;「数据在不在主机上」它也给了一个类型层面的选项,用 D 选。

动脑筋答案

TypedTensor<f64>                    280
TypedTensor<f64, Rank<2>>           216
TensorLayout<DynRank>               168
TensorLayout<Rank<2>>               40
ShapeVec                            80
  • 四个字段里只有 shape: R::Shape 随 R 变:ShapeVec 80 字节,[usize; 2] 16 字节,差 64,正好是 280 − 216。
  • 缓存的步长字段写死了 StrideVec,不是 R::Strides,所以 Rank<2> 的张量也带一个 80 字节的 TinyVec。TensorLayout<R> 的步长用的是 R::Strides,形状和步长都缩,才差 128。
  • 存储那一项也不随 R 变:Host 一支装的是 Vec<T>,与秩无关。静态秩的价值在类型检查和下标上,不在省内存上。

小结

  • tenferro-tensor-core 只剩元数据;HostTensor 删了,工作区里只有一套拥有型张量 TypedTensor<T, R, D>,tenferro_tensor::core 只转出元数据。
  • DType 由 define_scalar_tag! 的一张表展开(含 External(TypeId),24 字节);i32 和 f32 提升成 f64;直接运算不做隐式提升,混用 dtype 返回 dtype mismatch,eager 会提升。
  • R 是秩:DynRank(TinyVec 形状,秩 8 以内不上堆)或 Rank<N>([usize; N] 形状,可以 t[[i, j]])。
  • D 是表示:Host(一个 Vec,可 clone、可下标)、Gpu(分配组持有,不等于在 GPU 上)、Dynamic(默认,二者之一);转换不动数据,收窄失败把原值还回来。
  • Tensor 是七个 TypedTensor 的枚举加一个 External;as_typed 是安全的 Any 向下转型,into_typed 失败不再吞掉输入。
  • 大小:TypedTensor<f64> 280 字节(0.7.1 是 776),Tensor 312,视图 680;步长缓存(#2149)让建视图不分配。

💬 关联讨论

发起讨论
加载中...