如果你用惯了 PyTorch 或者 NumPy大概很难意识到 Tensor 本身其实不是一个“能装数据的列表”它更像一张精心设计的地图数据存放在某块连续内存里Tensor 只负责告诉你“从第几个字节开始每个维度跨多远取一次数”。最近我把手头一个叫 TensorPlay 的练习级框架重写了一遍核心目标就一个把 Tensor 背后的工程——从数据排布、内存分配到运算调度——全部自己实现一遍不依赖现成张量库来兜底。TensorPlay 的定位很朴素不追求超越 PyTorch而是用最克制的代码量把 Python 层常见的张量操作下沉到一块尊重内存布局的底层内核上。这个系列我会从张量本身讲起第一篇先聊三件事Tensor 的元数据构成、视图与副本的边界、运算内核是怎么被 dispatch 出来的。1. 在触碰存储之前Tensor 的三个元数据与它们如何共同定位一个数1.1 TensorPlay 想解决的真实痛点我在做 TensorPlay 之前有一段时间被 PyTorch 的“魔法”折腾得够呛。模型里某个 shape mismatch 的报错最终原因不是数据本身错了而是我对strides的理解不够一个看似正常的切片操作悄悄改变了内存访问模式后续的逐元素操作全部走了慢路径。更麻烦的是很多教程把 Tensor 描述成“多维数组”这个词会误导人——它让人以为 Tensor 内部真的是嵌套的 Python list于是你会下意识用“有几层列表”去理解维度一旦涉及到转置、切片、广播思路立刻断掉。TensorPlay 的第一版代码只有几千行但它强迫我面对一个核心问题如果抛开 PyTorch 的封装Tensor 至少要保存什么才能让任何一次索引操作都算得准答案比想象中简单也出乎意料地稳定一个原始字节缓冲区加上 shape、strides、storage_offset、dtype就构成了张量的全部物理信息。Tensor 不需要知道自己是“二维”还是“三维”它只需要知道在某个连续地址上每个逻辑索引应该偏移多少字节。1.2 从结构体开始理解 Tensor我用一个非常接近 C 侧设计的 Python 数据类来描述 TensorPlay 里最核心的对象from dataclasses import dataclass from typing import Tuple dataclass class Tensor: storage: bytearray # 真正的原始字节区域 shape: Tuple[int, ...] # 每个维度的长度 strides: Tuple[int, ...] # 每个维度前进 1 个元素时要跳多少个元素 offset: int # 从 storage 开头到第一个有效元素的距离 dtype: str # float32 / int64 ...给任何一个逻辑下标(i0, i1, i2)它在存储中的位置永远可以用一个公式算出来byte_pos dtype_size * (offset i0 * strides[0] i1 * strides[1] i2 * strides[2])这个公式是理解张量工程的钥匙。shape描述了“逻辑上这个张量有几行几列”strides描述了“物理上每走一步需要跨过多少个元素”。绝大多数疑难 bug最终都可以归约成这两者不一致。TensorPlay 里所有操作包括加法、切片、广播底层都在跟这个公式打交道性能优化的本质也只是想办法让这个公式在连续路径上退化成最简单的线性递增。很多初学者会问为什么不直接存一个嵌套数组因为嵌套数组有致命问题每个内层 list 是独立对象内存不一定连续缓存命中率极差更别提转置和切片时要复制整块数据。Tensor 用“元数据 连续缓冲区”的设计本质上是把“数据”和“数据形态”分离。修改 shape 只需要改元数据不需要碰字节这是张量库能高效处理大量数据的前提。1.3 为什么这三样信息已经足够有人可能觉得光有 shape 和 strides 不够还应该记录“这个张量从哪个父张量切出来的”。确实需要记录但记录方式就是offset。比如你想取一个 4x4 矩阵的中间 2x2 子块不需要新建存储只要在原有 storage 上把 offset 调整到子块起始位置再把 shape 改成(2, 2)strides 保持为(4, 1)一个轻量子视图就诞生了。TensorPlay 里所有 shape 变化比如reshape、transpose、broadcast_to第一版我都尝试用纯元数据变换去实现实在不行才分配新内存。这个约束逼着我大量使用 strides 去描述视图后来再看 PyTorch 的源码发现核心思路完全一致真正重量级的分配只发生在contiguous()、copy_()或者某些无法用步长表达的索引操作里。换句话说Tensor 的工程本质是“尽量改地图不搬运货物”。2. 内存分配与行主序为什么连续排布对算子性能如此重要2.1 连续内存不是天然存在的如果你把一个 Python list 传给 TensorPlay第一件事就是把它“压平”到一段连续内存里。为什么非压平不可因为 CPU 读取内存时不是按字节一个个读而是按缓存行cache line一次读 64 字节。如果数据在内存中连续那么你读前几个元素时后面几十个元素已经跟着进缓存了如果数据是东一块西一块的离散对象每次取数都要重新走一遍内存总线性能差距随数据量扩大变成数量级差异。TensorPlay 的 storage 默认用 64 字节对齐分配内存。这是个细节但直接决定后续能不能走 SIMD 优化路径。现代 CPU 的向量指令比如 AVX2一次能处理 256 位数据恰好等于 8 个 float32。如果内存基址不对齐加载指令可能要拆成两次甚至触发异常。所以 TensorPlay 里分配存储的函数不是直接调bytearray而是走一个自定义的对齐分配器保证每次分配的起始地址是 64 的倍数。2.2 行主序与步长的前世今生TensorPlay 默认采用行主序row-major也就是 C 语言的多维数组排布方式。一个 shape 为(2, 3)的 float32 张量内存里是这样排的索引 (0,0) (0,1) (0,2) | (1,0) (1,1) (1,2) 偏移 0 1 2 | 3 4 5对应的 strides 是(3, 1)。含义很直白第 0 维前进一个单位逻辑上跳到下一行物理上要跨过 3 个元素第 1 维前进一个单位跳到下一个元素物理上跨 1 个元素。行主序其实和“把多维数组按行拍平”的思路完全一致所以连续创建的张量最后一个维度的 stride 一定是 1。那有没有其他排布有列主序column-major就是反过来对应 strides(1, 2)。TensorPlay 不支持列主序但内置了对非连续视图的兼容。关键点在于非连续张量的 strides 不再是“从大到小”的规整序列。比如一个(3, 2)矩阵执行转置后得到 shape(2, 3)它的 strides 会变成(1, 3)意味着第二维反而要跳 3 个元素才能取到下一个数。此时数据在内存里仍然连续但逻辑上“隔三差五”取一次。我用一个小表格总结常见排布模式方便对照排查操作shapestrides是否连续新建 (2,3)(2,3)(3,1)是转置 (3,2).T(2,3)(1,3)否切片 [:, ::2](2,2)(6,2)否展平 reshape(-1)(6,)(1,)依赖于源是否连续2.3 内存池与引用计数小张量也有大开销TensorPlay 早期版本有个很明显的问题每个临时变量都分配独立 storage一次简单运算a b c会产生两次分配。虽然 Python 层对象会被 GC 回收但底层字节数组的释放和分配都涉及系统调用张量规模一大开销就很可观。解决思路是引入一个按 dtype 区分的内存池arena。TensorPlay 为每个 dtype 维护一个待回收存储块链表释放张量时不是真正销毁字节数组而是把 storage 放回池子。下次分配同 dtype、大小不超过阈值的张量时直接从池里取。这个策略对推理阶段特别有效因为推理时张量尺寸相对稳定内存池命中率很高。当然池化也会带来一个隐患如果某个超大张量长期占用池子内存不会立刻还给系统。TensorPlay 的处理是设定池容量上限超过上限的 storage 直接销毁避免静态内存无限膨胀。引用计数其实也藏在这里面。Python 层的__del__或weakref回调触发时机不可控所以 TensorPlay 没有把生命周期管理完全交给 Python 的 GC而是在 C 扩展层用一个简单的引用计数包裹 storage。基础规则是每次张量视图创建时对底层 storage 引用数加一视图销毁时减一减到零才允许 storage 回到内存池。这一步很容易出错如果你用原生bytearray做 storage视图和父张量的生命周期天然绑定看起来没事但一旦你想做“异步回收”或“跨线程共享数据”没有显式引用计数就会立刻翻车。2.4 连续性检测是一件必须做对的事底层内核执行任意操作前首先要问这个张量是不是连续内存TensorPlay 里有一个is_contiguous()检测从最后一个维度往前推最后一位 stride 必须等于 1倒数第二位 stride 必须等于倒数第一维长度乘以其 stride依次类推。这个检测非常便宜只是几个整数相乘和比较。它之所以重要是因为同一个算法可以针对连续和非连续两种内存布局分别写内核。连续路径可以用最简单的线性循环直接按顺序读内存非连续路径则必须保留 strides 逐点计算地址。TensorPlay 中同一操作通常有两版实现慢而兼容的非连续版本以及快而受限的连续版本。系统执行前先检测能走快速路径就绝不慢吞吞。很多深度学习框架里的contiguous()调用本质上就是在连续检测失败后重新分配一块连续内存并拷贝数据——这个操作是“兜底”而不是“默认”频繁触发它往往意味着你的内存访问模式有问题。3. 视图与副本哪些切片“免费”哪些切片暗藏复制3.1 视图的边界条件为什么不复制也能切出子块TensorPlay 里的基础切片规则和 Python 标准 list 切片长得一样但幕后完全不同。对a[1:3, 0:2]如果 a 是一个连续排布的(4, 4)矩阵切片结果是共享底层 storage 的视图。新张量的 shape 变成(2, 2)offset 指向a的(1, 0)元素strides 保持(4, 1)。整个过程没有发生任何数据拷贝所以切出一个上百万元素的子块也是常数级开销。这里有个容易忽略的事实子视图的 strides 通常继承自父张量而不是按新 shape 重新计算。很多人拿到一个切片后下意识认为它是连续的因为“它看起来像一个小矩阵”。实际上只有当你从父张量的第 0 个元素开始、按固定间距切到底并且没有跳步时子视图才可能保持连续。拿上面的(2, 2)子块来说它的 strides 是(4, 1)第二个维度 stride 还是 1但第一维跨度为 4说明下一行在内存里离上一行尾部还有两个被“跳过”的元素因此它不是连续的。如果你对这个切片做逐元素加法走慢速路径几乎是必然的。3.2 转置、跳步与最容易被忽视的广播视图除了基础切片TensorPlay 还有三种“免费”视图操作transpose交换 shape 和 strides 中对应位置的值。narrow/slice调整 offset 和 shapestrides 不变。expand/broadcast_to把某个维度长度为 1 的 stride 设为 0从而实现数据复用的逻辑扩展。第三种最反直觉。比如一个 shape 为(3, 1)的列向量要广播成(3, 4)参与矩阵运算TensorPlay 不会真的复制 4 份数据而是构造一个新视图它的 shape 是(3, 4)原维度的 stride 仍然是原来的行 stride新增维度的 stride 直接置 0。这样在第 1 维上不管取哪个下标寻址公式里这一项都等于 0永远访问同一个元素。数据没有复制视角却“变大”了。这种做法节省内存也带来了陷阱。如果你拿到一个广播视图然后想要原地修改它发现改动竟然同时影响“所有行”因为它们在内存里就是同一个位置。TensorPlay 的策略是有广播维度的张量默认设为只读任何写操作前先强制copy成真实独立数据。这一步很保守但能防止大量莫名其妙的数据串扰。3.3 花式索引为什么它总是让你“获得副本”切片规则可以用 strides 表示那么花式索引比如a[[0, 2], [1, 3]]还能用视图实现吗答案是基本不能。因为你提取的是两条不规则的坐标路径无法用单一 offset 和一组稳定的 strides 描述除非引入复杂的索引列表。TensorPlay 第一版对花式索引直接复制数据这样做省事也更接近大多数框架的语义。但这里藏着一个容易被误用的差异基础切片返回视图花式索引返回副本二者在后续赋值时行为完全不同。视图赋值会写回原张量副本赋值的改动则只停留在临时结果上。我排查过不止一个“数据怎么变了”的 bug最后都发现是有人把花式索引结果当成视图直接在上面做了原地修改。所以 TensorPlay 的文档里特意强调a[1:3]是视图a[[1, 2]]是副本二者之间的分界线可以简单记为“能否用 start/stop/step 描述”。3.4 调试张量布局的两板斧如果你也在写自己的张量内核我强烈建议打印两个东西strides 和连续性标志。我在 TensorPlay 里写了一个很小的辅助函数def debug_layout(t, nametensor): print(f{name}: shape{t.shape}, strides{t.strides}, foffset{t.offset}, contiguous{t.is_contiguous()})别小看这段代码。很多 shape 对不上、结果不对的问题只要打印出 strides立刻就能发现是转置后没连续化还是广播视图被误写。TensorPlay 内部也用这套信息生成调试日志当某个算子输出异常时直接把输入布局连同结果一起 dump 到文件省掉了反复加断点的痛苦。4. dtype 系统与运算调度字节要如何解释、算子如何找到实现4.1 dtype 是字节的解释契约不是类型标签TensorPlay 的 storage 是一个原始字节缓冲它本身不知道里面是整数还是浮点数。dtype字段的作用是告诉底层内核“每几个字节组成一个元素以及如何解释这串字节”。同一个四字节序列按 int32 看可能是 1065353216按 float32 看可能就是 1.0。这里我必须强调一个新手常犯的错误改变 dtype 不等于类型转换。如果你直接修改张量的 dtype 字段不做任何数据转换那就是在“重新解释”字节结果通常是灾难性的。TensorPlay 里区分了两个 APIt.view(dtypeint32) # 重新解释字节不改变底层数据 t.to(dtypefloat32) # 真正执行数值转换分配新存储前者是零拷贝的视图操作后者是重量级转换。很多人把这两个混为一谈最后得到一堆看起来像随机数的结果。工程上dtype 不仅仅是“类型标签”它还是调度表的关键索引不同 dtype 对应不同的内核实现加法的 float32 版本和 int64 版本不能混用。4.2 类型提升两个不同 dtype 相加谁迁就谁TensorPlay 里实现运算前的第一步是确定“输出 dtype”。规则参考了主流张量库布尔 整数 浮点数低精度向高精度提升有符号和无符号混合时取能同时容纳两者的类型。举个例子int32与float32相加结果是float32int8和int16相加结果是int16。这个决策看起来简单实现起来却繁琐因为要处理很多边界组合。TensorPlay 用一个两两组合的矩阵表来解决而不是一串 if-else。表越大越容易维护新增一个 dtype 时只要在矩阵中补一行一列即可。还有一个细节bool bool的结果不是 bool而是 int64因为 True True 要等于 2。如果输出 dtype 也用 bool数据直接溢出。这个规则在很多库里是统一的但第一次接触时很容易踩。我通常用一句口诀记住运算符的结果类型至少要能容纳所有参与运算的数值范围。4.3 dispatch 表避免写出一座 if-else 屎山TensorPlay 支持超过十种基础算子如果每种算子都按 dtype 写一遍分支代码会变成不可维护的多重嵌套判断。更好的做法是用字典做 dispatch 表。每个算子以(op_name, dtype, device)作为键对应的值是一个函数指针或可调用对象。dispatch_table { (add, float32, cpu): cpu_add_f32, (add, int64, cpu): cpu_add_i64, (mul, float32, cpu): cpu_mul_f32, # ... } def add(a, b): out_dtype promote_dtype(a.dtype, b.dtype) impl dispatch_table[(add, out_dtype, a.device)] return impl(a, b)这套设计的价值在于新增一个 dtype 时你不需要去改动几十个算子的函数体只需要为每个算子注册一个新的键值对。而且如果某个组合还没有实现查表会自然抛出 KeyError你能立刻收齐所有缺失的内核。TensorPlay 里我把 dispatch 表的注册和实现放在两个文件里头文件只声明函数签名具体实现按 dtype 拆分这样每个人都能快速定位自己负责的算子。实际运行中还有一层优化如果两个输入都是同一个 dtype直接查表如果不同先做类型提升然后把其中一个输入转换成输出 dtype再做同类型运算。和 NumPy 里的 upcast 逻辑类似只是我把转换时机提前到了算子入口避免每个内核算子都去判断“两个输入类型不同怎么办”。5. 广播在工程里的真实面孔零步长维度与通用循环5.1 广播规则的实质是“从右往左对齐”广播是新手最容易恐惧的概念之一但它的工程规则其实非常直白从最后一个维度开始往左对齐每个维度要么长度相同要么其中一个为 1要么某个张量在该维度不存在。不满足这些条件就直接报错。(3, 1)和(1, 4)可以广播成(3, 4)(3, 2)和(3,)也可以因为(3,)对齐到第二个张量的最后一维长度都是 3但(3, 2)和(2, 3)就不行因为对齐后第二维是 2 对 3谁也不等于谁且都不是 1。我见过很多人纠结“维度从哪里开始数”实际上只要记住“右对齐”三个字就够了。5.2 用零步长实现广播才是真正的工程做法你可能在网上看到过一些教程用np.broadcast_to然后说“它只是创建了一个视图”。但很少有人告诉你这个视图背后的机制是“维度扩展的 stride 被设置为 0”。TensorPlay 里专门有一个函数负责维度对齐前的 stride 适配def broadcast_strides(shape, target_shape): # 假设 target_shape 已经通过广播规则验证过 pad len(target_shape) - len(shape) padded_shape (1,) * pad shape strides [0] * len(target_shape) # 原张量在对应维度上的 stride 保留下来长度为 1 或新增维度的 stride 设为 0 for i, (s, ts) in enumerate(zip(padded_shape, target_shape)): strides[i] 0 if s 1 and ts ! 1 else original_stride(i) return tuple(strides)这个函数比较啰嗦但核心思想就一句目标 shape 上新增的维度或者原维度长度是 1 的维度都让 stride 等于 0。stride 0的含义是“寻址公式里这一项永远为零”于是该维度上任何索引都映射到同一个内存地址。5.3 通用广播循环让任意 shape 组合跑起来当我们把两个张量的 strides 都“适配”到目标输出 shape 后真正执行计算的循环可以写得很通用。TensorPlay 里最朴素的内核模板长这样def elementwise_broadcast_loop(out, a, b, op): out_shape out.shape for linear_idx in range(out.numel()): coords unravel_index(linear_idx, out_shape) a_off sum(coords[i] * a.strides[i] for i in range(len(out_shape))) b_off sum(coords[i] * b.strides[i] for i in range(len(out_shape))) out_buf[linear_idx] op(a_buf[a_off], b_buf[b_off])这个循环能处理任意 broadcast 组合代价是慢。每算一个输出元素都要做两次多维坐标的逐维乘加。所以 TensorPlay 在通用循环之前会先检查一个特例两个输入是否都是连续且 shape 完全相同。如果是直接用一重线性循环把 strides 全部忽略速度提升几十倍。这个特例覆盖了深度学习里绝大多数逐元素算子场景。工程上我倾向这样理解广播性能如果输出张量是连续的且至少一个输入不是连续或 shape 不同那广播循环的瓶颈在地址计算但如果输出张量本身非连续那内存访问模式已经决定了性能地址计算反而是次要矛盾。先解决连续性问题再谈广播优化顺序不要反。5.4 广播报错的艺术信息要能帮你定位维度最后说一个 TensorPlay 开发中很实在的经验广播报错信息一定要把两个原始 shape 和目标 shape 一起打印最好指明从哪一位开始冲突。ValueError: 无法广播: 左张量形状 (3, 2)右张量形状 (2, 3) 在从右往左的第 2 个维度冲突: 2 ! 3这样的报错能让你少 debug 十分钟。早期我只写“shape mismatch”结果排查的人完全不知道是哪个维度出问题。后来改成逐步对齐检测每检查一个维度就记录冲突位置再拼接成详细错误。如果你也在写类似框架强烈建议现在就把这条改掉。我在实际写 TensorPlay 的过程中最大的体会是Tensor 背后没有玄学。dtype 决定字节如何解释strides 决定地址如何换算storage 决定数据住在哪里把这三条主线刻在脑子里绝大多数张量工程问题都能沿着这条线索一步步拆下去。这个系列后续还会聊算子内核的向量化优化、自动求导如何追踪图、以及多设备的同步问题但第一步永远是先把数据布局这只大象摸清楚。
深入拆解 Tensor 工程：内存布局、视图与广播的实现
如果你用惯了 PyTorch 或者 NumPy大概很难意识到 Tensor 本身其实不是一个“能装数据的列表”它更像一张精心设计的地图数据存放在某块连续内存里Tensor 只负责告诉你“从第几个字节开始每个维度跨多远取一次数”。最近我把手头一个叫 TensorPlay 的练习级框架重写了一遍核心目标就一个把 Tensor 背后的工程——从数据排布、内存分配到运算调度——全部自己实现一遍不依赖现成张量库来兜底。TensorPlay 的定位很朴素不追求超越 PyTorch而是用最克制的代码量把 Python 层常见的张量操作下沉到一块尊重内存布局的底层内核上。这个系列我会从张量本身讲起第一篇先聊三件事Tensor 的元数据构成、视图与副本的边界、运算内核是怎么被 dispatch 出来的。1. 在触碰存储之前Tensor 的三个元数据与它们如何共同定位一个数1.1 TensorPlay 想解决的真实痛点我在做 TensorPlay 之前有一段时间被 PyTorch 的“魔法”折腾得够呛。模型里某个 shape mismatch 的报错最终原因不是数据本身错了而是我对strides的理解不够一个看似正常的切片操作悄悄改变了内存访问模式后续的逐元素操作全部走了慢路径。更麻烦的是很多教程把 Tensor 描述成“多维数组”这个词会误导人——它让人以为 Tensor 内部真的是嵌套的 Python list于是你会下意识用“有几层列表”去理解维度一旦涉及到转置、切片、广播思路立刻断掉。TensorPlay 的第一版代码只有几千行但它强迫我面对一个核心问题如果抛开 PyTorch 的封装Tensor 至少要保存什么才能让任何一次索引操作都算得准答案比想象中简单也出乎意料地稳定一个原始字节缓冲区加上 shape、strides、storage_offset、dtype就构成了张量的全部物理信息。Tensor 不需要知道自己是“二维”还是“三维”它只需要知道在某个连续地址上每个逻辑索引应该偏移多少字节。1.2 从结构体开始理解 Tensor我用一个非常接近 C 侧设计的 Python 数据类来描述 TensorPlay 里最核心的对象from dataclasses import dataclass from typing import Tuple dataclass class Tensor: storage: bytearray # 真正的原始字节区域 shape: Tuple[int, ...] # 每个维度的长度 strides: Tuple[int, ...] # 每个维度前进 1 个元素时要跳多少个元素 offset: int # 从 storage 开头到第一个有效元素的距离 dtype: str # float32 / int64 ...给任何一个逻辑下标(i0, i1, i2)它在存储中的位置永远可以用一个公式算出来byte_pos dtype_size * (offset i0 * strides[0] i1 * strides[1] i2 * strides[2])这个公式是理解张量工程的钥匙。shape描述了“逻辑上这个张量有几行几列”strides描述了“物理上每走一步需要跨过多少个元素”。绝大多数疑难 bug最终都可以归约成这两者不一致。TensorPlay 里所有操作包括加法、切片、广播底层都在跟这个公式打交道性能优化的本质也只是想办法让这个公式在连续路径上退化成最简单的线性递增。很多初学者会问为什么不直接存一个嵌套数组因为嵌套数组有致命问题每个内层 list 是独立对象内存不一定连续缓存命中率极差更别提转置和切片时要复制整块数据。Tensor 用“元数据 连续缓冲区”的设计本质上是把“数据”和“数据形态”分离。修改 shape 只需要改元数据不需要碰字节这是张量库能高效处理大量数据的前提。1.3 为什么这三样信息已经足够有人可能觉得光有 shape 和 strides 不够还应该记录“这个张量从哪个父张量切出来的”。确实需要记录但记录方式就是offset。比如你想取一个 4x4 矩阵的中间 2x2 子块不需要新建存储只要在原有 storage 上把 offset 调整到子块起始位置再把 shape 改成(2, 2)strides 保持为(4, 1)一个轻量子视图就诞生了。TensorPlay 里所有 shape 变化比如reshape、transpose、broadcast_to第一版我都尝试用纯元数据变换去实现实在不行才分配新内存。这个约束逼着我大量使用 strides 去描述视图后来再看 PyTorch 的源码发现核心思路完全一致真正重量级的分配只发生在contiguous()、copy_()或者某些无法用步长表达的索引操作里。换句话说Tensor 的工程本质是“尽量改地图不搬运货物”。2. 内存分配与行主序为什么连续排布对算子性能如此重要2.1 连续内存不是天然存在的如果你把一个 Python list 传给 TensorPlay第一件事就是把它“压平”到一段连续内存里。为什么非压平不可因为 CPU 读取内存时不是按字节一个个读而是按缓存行cache line一次读 64 字节。如果数据在内存中连续那么你读前几个元素时后面几十个元素已经跟着进缓存了如果数据是东一块西一块的离散对象每次取数都要重新走一遍内存总线性能差距随数据量扩大变成数量级差异。TensorPlay 的 storage 默认用 64 字节对齐分配内存。这是个细节但直接决定后续能不能走 SIMD 优化路径。现代 CPU 的向量指令比如 AVX2一次能处理 256 位数据恰好等于 8 个 float32。如果内存基址不对齐加载指令可能要拆成两次甚至触发异常。所以 TensorPlay 里分配存储的函数不是直接调bytearray而是走一个自定义的对齐分配器保证每次分配的起始地址是 64 的倍数。2.2 行主序与步长的前世今生TensorPlay 默认采用行主序row-major也就是 C 语言的多维数组排布方式。一个 shape 为(2, 3)的 float32 张量内存里是这样排的索引 (0,0) (0,1) (0,2) | (1,0) (1,1) (1,2) 偏移 0 1 2 | 3 4 5对应的 strides 是(3, 1)。含义很直白第 0 维前进一个单位逻辑上跳到下一行物理上要跨过 3 个元素第 1 维前进一个单位跳到下一个元素物理上跨 1 个元素。行主序其实和“把多维数组按行拍平”的思路完全一致所以连续创建的张量最后一个维度的 stride 一定是 1。那有没有其他排布有列主序column-major就是反过来对应 strides(1, 2)。TensorPlay 不支持列主序但内置了对非连续视图的兼容。关键点在于非连续张量的 strides 不再是“从大到小”的规整序列。比如一个(3, 2)矩阵执行转置后得到 shape(2, 3)它的 strides 会变成(1, 3)意味着第二维反而要跳 3 个元素才能取到下一个数。此时数据在内存里仍然连续但逻辑上“隔三差五”取一次。我用一个小表格总结常见排布模式方便对照排查操作shapestrides是否连续新建 (2,3)(2,3)(3,1)是转置 (3,2).T(2,3)(1,3)否切片 [:, ::2](2,2)(6,2)否展平 reshape(-1)(6,)(1,)依赖于源是否连续2.3 内存池与引用计数小张量也有大开销TensorPlay 早期版本有个很明显的问题每个临时变量都分配独立 storage一次简单运算a b c会产生两次分配。虽然 Python 层对象会被 GC 回收但底层字节数组的释放和分配都涉及系统调用张量规模一大开销就很可观。解决思路是引入一个按 dtype 区分的内存池arena。TensorPlay 为每个 dtype 维护一个待回收存储块链表释放张量时不是真正销毁字节数组而是把 storage 放回池子。下次分配同 dtype、大小不超过阈值的张量时直接从池里取。这个策略对推理阶段特别有效因为推理时张量尺寸相对稳定内存池命中率很高。当然池化也会带来一个隐患如果某个超大张量长期占用池子内存不会立刻还给系统。TensorPlay 的处理是设定池容量上限超过上限的 storage 直接销毁避免静态内存无限膨胀。引用计数其实也藏在这里面。Python 层的__del__或weakref回调触发时机不可控所以 TensorPlay 没有把生命周期管理完全交给 Python 的 GC而是在 C 扩展层用一个简单的引用计数包裹 storage。基础规则是每次张量视图创建时对底层 storage 引用数加一视图销毁时减一减到零才允许 storage 回到内存池。这一步很容易出错如果你用原生bytearray做 storage视图和父张量的生命周期天然绑定看起来没事但一旦你想做“异步回收”或“跨线程共享数据”没有显式引用计数就会立刻翻车。2.4 连续性检测是一件必须做对的事底层内核执行任意操作前首先要问这个张量是不是连续内存TensorPlay 里有一个is_contiguous()检测从最后一个维度往前推最后一位 stride 必须等于 1倒数第二位 stride 必须等于倒数第一维长度乘以其 stride依次类推。这个检测非常便宜只是几个整数相乘和比较。它之所以重要是因为同一个算法可以针对连续和非连续两种内存布局分别写内核。连续路径可以用最简单的线性循环直接按顺序读内存非连续路径则必须保留 strides 逐点计算地址。TensorPlay 中同一操作通常有两版实现慢而兼容的非连续版本以及快而受限的连续版本。系统执行前先检测能走快速路径就绝不慢吞吞。很多深度学习框架里的contiguous()调用本质上就是在连续检测失败后重新分配一块连续内存并拷贝数据——这个操作是“兜底”而不是“默认”频繁触发它往往意味着你的内存访问模式有问题。3. 视图与副本哪些切片“免费”哪些切片暗藏复制3.1 视图的边界条件为什么不复制也能切出子块TensorPlay 里的基础切片规则和 Python 标准 list 切片长得一样但幕后完全不同。对a[1:3, 0:2]如果 a 是一个连续排布的(4, 4)矩阵切片结果是共享底层 storage 的视图。新张量的 shape 变成(2, 2)offset 指向a的(1, 0)元素strides 保持(4, 1)。整个过程没有发生任何数据拷贝所以切出一个上百万元素的子块也是常数级开销。这里有个容易忽略的事实子视图的 strides 通常继承自父张量而不是按新 shape 重新计算。很多人拿到一个切片后下意识认为它是连续的因为“它看起来像一个小矩阵”。实际上只有当你从父张量的第 0 个元素开始、按固定间距切到底并且没有跳步时子视图才可能保持连续。拿上面的(2, 2)子块来说它的 strides 是(4, 1)第二个维度 stride 还是 1但第一维跨度为 4说明下一行在内存里离上一行尾部还有两个被“跳过”的元素因此它不是连续的。如果你对这个切片做逐元素加法走慢速路径几乎是必然的。3.2 转置、跳步与最容易被忽视的广播视图除了基础切片TensorPlay 还有三种“免费”视图操作transpose交换 shape 和 strides 中对应位置的值。narrow/slice调整 offset 和 shapestrides 不变。expand/broadcast_to把某个维度长度为 1 的 stride 设为 0从而实现数据复用的逻辑扩展。第三种最反直觉。比如一个 shape 为(3, 1)的列向量要广播成(3, 4)参与矩阵运算TensorPlay 不会真的复制 4 份数据而是构造一个新视图它的 shape 是(3, 4)原维度的 stride 仍然是原来的行 stride新增维度的 stride 直接置 0。这样在第 1 维上不管取哪个下标寻址公式里这一项都等于 0永远访问同一个元素。数据没有复制视角却“变大”了。这种做法节省内存也带来了陷阱。如果你拿到一个广播视图然后想要原地修改它发现改动竟然同时影响“所有行”因为它们在内存里就是同一个位置。TensorPlay 的策略是有广播维度的张量默认设为只读任何写操作前先强制copy成真实独立数据。这一步很保守但能防止大量莫名其妙的数据串扰。3.3 花式索引为什么它总是让你“获得副本”切片规则可以用 strides 表示那么花式索引比如a[[0, 2], [1, 3]]还能用视图实现吗答案是基本不能。因为你提取的是两条不规则的坐标路径无法用单一 offset 和一组稳定的 strides 描述除非引入复杂的索引列表。TensorPlay 第一版对花式索引直接复制数据这样做省事也更接近大多数框架的语义。但这里藏着一个容易被误用的差异基础切片返回视图花式索引返回副本二者在后续赋值时行为完全不同。视图赋值会写回原张量副本赋值的改动则只停留在临时结果上。我排查过不止一个“数据怎么变了”的 bug最后都发现是有人把花式索引结果当成视图直接在上面做了原地修改。所以 TensorPlay 的文档里特意强调a[1:3]是视图a[[1, 2]]是副本二者之间的分界线可以简单记为“能否用 start/stop/step 描述”。3.4 调试张量布局的两板斧如果你也在写自己的张量内核我强烈建议打印两个东西strides 和连续性标志。我在 TensorPlay 里写了一个很小的辅助函数def debug_layout(t, nametensor): print(f{name}: shape{t.shape}, strides{t.strides}, foffset{t.offset}, contiguous{t.is_contiguous()})别小看这段代码。很多 shape 对不上、结果不对的问题只要打印出 strides立刻就能发现是转置后没连续化还是广播视图被误写。TensorPlay 内部也用这套信息生成调试日志当某个算子输出异常时直接把输入布局连同结果一起 dump 到文件省掉了反复加断点的痛苦。4. dtype 系统与运算调度字节要如何解释、算子如何找到实现4.1 dtype 是字节的解释契约不是类型标签TensorPlay 的 storage 是一个原始字节缓冲它本身不知道里面是整数还是浮点数。dtype字段的作用是告诉底层内核“每几个字节组成一个元素以及如何解释这串字节”。同一个四字节序列按 int32 看可能是 1065353216按 float32 看可能就是 1.0。这里我必须强调一个新手常犯的错误改变 dtype 不等于类型转换。如果你直接修改张量的 dtype 字段不做任何数据转换那就是在“重新解释”字节结果通常是灾难性的。TensorPlay 里区分了两个 APIt.view(dtypeint32) # 重新解释字节不改变底层数据 t.to(dtypefloat32) # 真正执行数值转换分配新存储前者是零拷贝的视图操作后者是重量级转换。很多人把这两个混为一谈最后得到一堆看起来像随机数的结果。工程上dtype 不仅仅是“类型标签”它还是调度表的关键索引不同 dtype 对应不同的内核实现加法的 float32 版本和 int64 版本不能混用。4.2 类型提升两个不同 dtype 相加谁迁就谁TensorPlay 里实现运算前的第一步是确定“输出 dtype”。规则参考了主流张量库布尔 整数 浮点数低精度向高精度提升有符号和无符号混合时取能同时容纳两者的类型。举个例子int32与float32相加结果是float32int8和int16相加结果是int16。这个决策看起来简单实现起来却繁琐因为要处理很多边界组合。TensorPlay 用一个两两组合的矩阵表来解决而不是一串 if-else。表越大越容易维护新增一个 dtype 时只要在矩阵中补一行一列即可。还有一个细节bool bool的结果不是 bool而是 int64因为 True True 要等于 2。如果输出 dtype 也用 bool数据直接溢出。这个规则在很多库里是统一的但第一次接触时很容易踩。我通常用一句口诀记住运算符的结果类型至少要能容纳所有参与运算的数值范围。4.3 dispatch 表避免写出一座 if-else 屎山TensorPlay 支持超过十种基础算子如果每种算子都按 dtype 写一遍分支代码会变成不可维护的多重嵌套判断。更好的做法是用字典做 dispatch 表。每个算子以(op_name, dtype, device)作为键对应的值是一个函数指针或可调用对象。dispatch_table { (add, float32, cpu): cpu_add_f32, (add, int64, cpu): cpu_add_i64, (mul, float32, cpu): cpu_mul_f32, # ... } def add(a, b): out_dtype promote_dtype(a.dtype, b.dtype) impl dispatch_table[(add, out_dtype, a.device)] return impl(a, b)这套设计的价值在于新增一个 dtype 时你不需要去改动几十个算子的函数体只需要为每个算子注册一个新的键值对。而且如果某个组合还没有实现查表会自然抛出 KeyError你能立刻收齐所有缺失的内核。TensorPlay 里我把 dispatch 表的注册和实现放在两个文件里头文件只声明函数签名具体实现按 dtype 拆分这样每个人都能快速定位自己负责的算子。实际运行中还有一层优化如果两个输入都是同一个 dtype直接查表如果不同先做类型提升然后把其中一个输入转换成输出 dtype再做同类型运算。和 NumPy 里的 upcast 逻辑类似只是我把转换时机提前到了算子入口避免每个内核算子都去判断“两个输入类型不同怎么办”。5. 广播在工程里的真实面孔零步长维度与通用循环5.1 广播规则的实质是“从右往左对齐”广播是新手最容易恐惧的概念之一但它的工程规则其实非常直白从最后一个维度开始往左对齐每个维度要么长度相同要么其中一个为 1要么某个张量在该维度不存在。不满足这些条件就直接报错。(3, 1)和(1, 4)可以广播成(3, 4)(3, 2)和(3,)也可以因为(3,)对齐到第二个张量的最后一维长度都是 3但(3, 2)和(2, 3)就不行因为对齐后第二维是 2 对 3谁也不等于谁且都不是 1。我见过很多人纠结“维度从哪里开始数”实际上只要记住“右对齐”三个字就够了。5.2 用零步长实现广播才是真正的工程做法你可能在网上看到过一些教程用np.broadcast_to然后说“它只是创建了一个视图”。但很少有人告诉你这个视图背后的机制是“维度扩展的 stride 被设置为 0”。TensorPlay 里专门有一个函数负责维度对齐前的 stride 适配def broadcast_strides(shape, target_shape): # 假设 target_shape 已经通过广播规则验证过 pad len(target_shape) - len(shape) padded_shape (1,) * pad shape strides [0] * len(target_shape) # 原张量在对应维度上的 stride 保留下来长度为 1 或新增维度的 stride 设为 0 for i, (s, ts) in enumerate(zip(padded_shape, target_shape)): strides[i] 0 if s 1 and ts ! 1 else original_stride(i) return tuple(strides)这个函数比较啰嗦但核心思想就一句目标 shape 上新增的维度或者原维度长度是 1 的维度都让 stride 等于 0。stride 0的含义是“寻址公式里这一项永远为零”于是该维度上任何索引都映射到同一个内存地址。5.3 通用广播循环让任意 shape 组合跑起来当我们把两个张量的 strides 都“适配”到目标输出 shape 后真正执行计算的循环可以写得很通用。TensorPlay 里最朴素的内核模板长这样def elementwise_broadcast_loop(out, a, b, op): out_shape out.shape for linear_idx in range(out.numel()): coords unravel_index(linear_idx, out_shape) a_off sum(coords[i] * a.strides[i] for i in range(len(out_shape))) b_off sum(coords[i] * b.strides[i] for i in range(len(out_shape))) out_buf[linear_idx] op(a_buf[a_off], b_buf[b_off])这个循环能处理任意 broadcast 组合代价是慢。每算一个输出元素都要做两次多维坐标的逐维乘加。所以 TensorPlay 在通用循环之前会先检查一个特例两个输入是否都是连续且 shape 完全相同。如果是直接用一重线性循环把 strides 全部忽略速度提升几十倍。这个特例覆盖了深度学习里绝大多数逐元素算子场景。工程上我倾向这样理解广播性能如果输出张量是连续的且至少一个输入不是连续或 shape 不同那广播循环的瓶颈在地址计算但如果输出张量本身非连续那内存访问模式已经决定了性能地址计算反而是次要矛盾。先解决连续性问题再谈广播优化顺序不要反。5.4 广播报错的艺术信息要能帮你定位维度最后说一个 TensorPlay 开发中很实在的经验广播报错信息一定要把两个原始 shape 和目标 shape 一起打印最好指明从哪一位开始冲突。ValueError: 无法广播: 左张量形状 (3, 2)右张量形状 (2, 3) 在从右往左的第 2 个维度冲突: 2 ! 3这样的报错能让你少 debug 十分钟。早期我只写“shape mismatch”结果排查的人完全不知道是哪个维度出问题。后来改成逐步对齐检测每检查一个维度就记录冲突位置再拼接成详细错误。如果你也在写类似框架强烈建议现在就把这条改掉。我在实际写 TensorPlay 的过程中最大的体会是Tensor 背后没有玄学。dtype 决定字节如何解释strides 决定地址如何换算storage 决定数据住在哪里把这三条主线刻在脑子里绝大多数张量工程问题都能沿着这条线索一步步拆下去。这个系列后续还会聊算子内核的向量化优化、自动求导如何追踪图、以及多设备的同步问题但第一步永远是先把数据布局这只大象摸清楚。
