第2章 Tensor、Storage、TensorImpl 三件套
“A tensor is just a multi-dimensional array. Until you start asking how it shares memory with other tensors.”
—— Edward Yang, “PyTorch internals” (2019)
本章要点
- 用户看到的
torch.Tensor背后是 三层 C++ 数据结构:TensorBase→TensorImpl→Storage(+StorageImpl) view不复制内存:因为它共享同一个Storage,只复制TensorImpl元信息(sizes / strides / offset)- strides 与 offset 是张量的”指针计算配方”:每一次索引
t[i,j,k]都翻译成data + (i·s0 + j·s1 + k·s2 + offset) × dtype_size TensorImpl是 PyTorch 内存最敏感的类之一,它手动用位域压到约 192 字节,因为每秒可能创建上百万次is_contiguous_是缓存属性而非动态计算,PyTorch 在每次修改 strides 时手动同步这一缓存reshape≈view if contiguous else contiguous().view:理解三件套就理解为什么有时reshape会偷偷复制
2.1 一个 view 的悖论
观察这段代码:
import torch
a = torch.arange(12) # [0, 1, 2, ..., 11]
b = a.view(3, 4) # 视图为 [[0,1,2,3], [4,5,6,7], [8,9,10,11]]
b[0, 0] = 999
print(a) # 输出: [999, 1, 2, ..., 11]
print(a.data_ptr() == b.data_ptr()) # True
b 是一个新的 Python 对象,shape 完全不同([12] vs [3, 4])。但修改 b 居然影响了 a,且两者的 data_ptr 指向同一段内存。
更奇特的是:
c = a.view(4, 3) # 又一个新视图
print(a.data_ptr() == c.data_ptr()) # True
print(b.data_ptr() == c.data_ptr()) # True
# 三个张量共享同一段 12 个 int64 的物理内存
如果你只把张量想成”数组+形状”,这件事是反直觉的:怎么可能 12 个字节同时是三个不同形状的张量?
答案要从 PyTorch 的 三件套 数据结构说起。从这一章开始,我们正式离开”用户视角”,进入 PyTorch 内部的 C++ 世界。
2.2 三件套全景图
打开 PyTorch 源码,你会发现一个张量在 C++ 里其实是三层嵌套的对象:
graph TB
subgraph Python["Python 用户层"]
T[torch.Tensor 实例]
end
subgraph CppFrontend["C++ 前端层 (按值传递的轻包装)"]
TB[at::TensorBase / at::Tensor<br/>只有一个 intrusive_ptr 字段]
end
subgraph CppCore["C++ 核心层 (引用计数管理)"]
TI["c10::TensorImpl<br/>──────────<br/>Storage storage_<br/>SizesAndStrides sizes_and_strides_<br/>storage_offset_ / numel_<br/>data_type_ / device_opt_<br/>key_set_ / autograd_meta_<br/>位域 (~20 个 1-bit 标志)"]
end
subgraph Storage["存储层 (实际数据)"]
ST[c10::Storage<br/>持有 intrusive_ptr<StorageImpl>]
SI["c10::StorageImpl<br/>──────────<br/>DataPtr data_ptr_<br/>SymInt size_bytes_<br/>Allocator* allocator_<br/>refcount + weakcount"]
end
T -. CPython slot .-> TB
TB -- intrusive_ptr --> TI
TI -- 嵌入对象 --> ST
ST -- intrusive_ptr --> SI
SI -. 拥有 .-> RAW["实际的 GPU/CPU 显存<br/>byte buffer"]
style TI fill:#fef3c7,stroke:#f59e0b,stroke-width:2px
style SI fill:#dbeafe,stroke:#3b82f6,stroke-width:2px
style RAW fill:#dcfce7,stroke:#22c55e,stroke-width:2px
这张图的关键信息:
Tensor/TensorBase只是一个 指针包装类,里面只有一个intrusive_ptr<TensorImpl>字段。复制Tensor几乎免费 —— 只是 +1 引用计数TensorImpl是真正”长成张量样子”的对象 —— 形状、步长、dtype、device 这些元信息都在它身上Storage/StorageImpl又是另一层 —— 它只管 底层的字节缓冲区,对张量的形状一无所知
view 共享内存的秘密就在这里:a.view(3, 4) 创建一个新的 TensorImpl,但新 TensorImpl 的 storage_ 字段和老 TensorImpl 共享同一个 StorageImpl。两个 TensorImpl 各自有自己的 sizes/strides/offset,但都指向同一块内存。
下面我们逐层拆开。
2.3 第一层:TensorBase —— 一个指针的轻包装
PyTorch 用户用的 torch.Tensor 在 C++ 里对应的核心类是 at::TensorBase,定义在:
// aten/src/ATen/core/TensorBase.h:95
class TORCH_API TensorBase {
public:
...
protected:
c10::intrusive_ptr<TensorImpl, UndefinedTensorImpl> impl_;
};
整个 TensorBase 只有一个数据成员:impl_,类型是 intrusive_ptr<TensorImpl> —— 一个智能指针。
为什么这么薄?因为 PyTorch 中 Tensor 经常按值在函数间传递(Tensor add(Tensor a, Tensor b))。如果 Tensor 本身很重,按值传递的开销就高。把所有真正的数据放到 heap 上的 TensorImpl,让 Tensor 只是一个 8 字节指针 —— 传递成本和传递一个 int64_t 没差。
at::Tensor 是 TensorBase 的子类,在 aten/src/ATen/core/Tensor.h 里被代码生成器扩展出几千个方法(Tensor::add / Tensor::matmul / …)。但这些方法本质都是查 impl_ 的字段或者把请求转发给 dispatcher。
2.4 第二层:TensorImpl —— 张量真正的”骨架”
TensorImpl 是 PyTorch 最重要的 C++ 类之一,定义在 c10/core/TensorImpl.h:511:
struct C10_API TensorImpl : public c10::intrusive_ptr_target {
protected:
Storage storage_; // 指向数据的 Storage
private:
std::unique_ptr<AutogradMetaInterface> autograd_meta_ = nullptr;
protected:
std::unique_ptr<c10::ExtraMeta> extra_meta_ = nullptr;
c10::VariableVersion version_counter_; // autograd inplace 检测
impl::PyObjectSlot pyobj_slot_; // Python 对象绑定
c10::impl::SizesAndStrides sizes_and_strides_; // 形状与步长
int64_t storage_offset_ = 0; // 在 storage 内的起点
int64_t numel_ = 1; // 元素总数
caffe2::TypeMeta data_type_; // dtype
std::optional<c10::Device> device_opt_; // device
// 大量位域:is_contiguous_, is_channels_last_, allow_metadata_change_, ...
...
private:
DispatchKeySet key_set_; // 第 1 章 §1.3 提过
};
——c10/core/TensorImpl.h 的字段(精简)
注意三件事:
2.4.1 它继承了 intrusive_ptr_target
TensorImpl 不能用 std::shared_ptr 管 —— 它继承的是 PyTorch 自己的 intrusive_ptr_target,里面藏着一个 std::atomic<uint64_t> combined_refcount_(c10/util/intrusive_ptr.h:188),把 strong refcount 和 weak refcount 打包成一个 64-bit 原子量。第 3 章会专门讲为什么 PyTorch 不用 shared_ptr —— 简单说:省一次堆分配 + 把引用计数贴在对象本身上能让 cache 更友好。
2.4.2 它的字段被极致压缩
仔细数一下 TensorImpl 在 v2.11 里的字段:
| 字段 | 大小 |
|---|---|
Storage storage_ | 8 字节(一个 intrusive_ptr) |
autograd_meta_ | 8 字节 |
extra_meta_ | 8 字节 |
version_counter_ | 8 字节 |
pyobj_slot_ | 16 字节 |
sizes_and_strides_ | 88 字节(小张量内联,长张量堆分配) |
storage_offset_ | 8 字节 |
numel_ | 8 字节 |
data_type_ | 2 字节 |
device_opt_ | 4 字节 |
key_set_ | 8 字节 |
| 位域(约 20 个 1-bit 标志) | ~3 字节 |
combined_refcount_(来自基类) | 8 字节 |
加起来约 180 字节。看起来不大?但乘上一次大模型训练每秒可能创建几十万次张量 —— 一秒钟就是几十兆字节的纯元数据。
PyTorch 团队为这个开销做了很多苦工:
SizesAndStrides对 ≤ 5 维张量不堆分配(直接 inline),覆盖 95%+ 真实场景- 大量 bool 字段都用
:1位域压成 1 bit,不是一个 1 字节的 bool device_opt_用std::optional<Device>而不是Device*,省一次堆分配autograd_meta_默认nullptr,只在需要梯度时分配(这是为什么 inference 模式比训练模式快不少)
理解 TensorImpl 的字段布局,你就理解了为什么 PyTorch 团队对”加一个新字段”这件事极其谨慎 —— 每多一字节,都要乘上每秒几十万的张量创建。
打开 c10/core/TensorImpl.h 你会看到字段顺序排得格外讲究 —— 引用计数和最常用的字段(storage_、sizes_and_strides_)排在前面,让常用属性命中同一条 cache line;位域被集中到一段连续区域便于编译器一次取出整个字节。这是 C++ 工程里典型的 “先按访问频率排,再按对齐塞填” 的优化手法。普通业务代码没必要这样做,但对一个每秒被实例化几十万次的核心类,这种”卡尺级别”的字段排布就是性能的真正来源。
注释里有一段 Edward Yang 留下的警句(c10/core/TensorImpl.h:444-470 附近):每加一个字段前先想清楚它能不能放进 extra_meta_ 这个”溢出区”。extra_meta_ 是一个 std::unique_ptr,平时是 nullptr,只在张量真的需要某些罕见元信息(如 namedtensor、symbolic shapes)时才分配。这种 “按需付费” 的内存策略让常态张量保持苗条,又不堵死扩展路径。
2.4.3 sizes_and_strides_:形状与步长的双胞胎容器
SizesAndStrides(c10/core/impl/SizesAndStrides.h)是一个非常巧妙的容器,它同时存储 sizes 和 strides 两个数组。设计思路:
- 对维度数 ≤ 5 的张量(绝大多数):直接在对象内联两个
int64_t[5],零堆分配 - 对维度数 > 5 的张量:退化到堆分配的
SmallVector
为什么把 sizes 和 strides 绑在一起?因为它们永远成对出现,一起读、一起写。把它们放在同一个连续内存块里,CPU 的 cache line 一次就能拿到两个 —— 这是 PyTorch 数据结构设计的典型权衡。
值得一看的是 SizesAndStrides::resize 的源码(c10/core/impl/SizesAndStrides.h):当从 inline 模式(≤5 维)跨过 5 维边界时,它会一次性 malloc 一段 2 * dim * sizeof(int64_t) 的内存,把 inline 部分的数据整体拷过去,然后切换到 outline 模式。回头从 outline 缩回 inline 时反向操作。这种”双模式存储”在 C++ 标准库里能找到对应物 —— std::string 的 SSO(Small String Optimization)就是同一种思想。理解这种 “小数据走快车道、大数据走慢车道” 的模式,你就能在写自己的高性能容器时直接借鉴。
2.5 第三层:Storage 与 StorageImpl —— 真正的数据
TensorImpl 通过 storage_ 字段持有一个 Storage(按值嵌入):
// c10/core/Storage.h:25
struct C10_API Storage {
...
protected:
c10::intrusive_ptr<StorageImpl> storage_impl_;
};
Storage 又是一个轻包装,里面只有一个 intrusive_ptr<StorageImpl>。真正的数据在 StorageImpl:
// c10/core/StorageImpl.h:52
struct C10_API StorageImpl : public c10::intrusive_ptr_target {
...
private:
DataPtr data_ptr_; // 数据指针 + 析构器 + device
SymInt size_bytes_; // 字节数(支持符号大小)
bool size_bytes_is_heap_allocated_;
bool resizable_;
bool received_cuda_;
Allocator* allocator_; // 分配器(CUDA Caching Allocator 在这里)
...
};
——c10/core/StorageImpl.h:52-100(精简)
StorageImpl 的核心字段是 data_ptr_:一个 c10::DataPtr,包含:
- 实际的数据指针(
void*) - device 信息
- 析构器(在自己被销毁时怎么 free 这块内存)
正是 data_ptr_ 这个析构器机制,让 PyTorch 能把 CUDA Caching Allocator、CPU malloc、Pinned Memory、共享内存等多种内存源统一抽象成一个 Storage。第 4 章会详细讲。
2.5.0 typed vs untyped storage:v1.x → v2.x 的 dtype 解耦
源码里有一段值得介绍的演进。在 PyTorch v1.x 里,Storage 是 类型化 的 —— FloatStorage / DoubleStorage / IntStorage 是不同的 C++ 类型,分别绑定到 fp32 / fp64 / int32 等 dtype。这个设计有问题:同一段物理内存如果用户想”reinterpret”成另一种 dtype,必须真的复制一份数据。
v1.5 之后 PyTorch 团队把 storage 重构为 untyped —— 现在的 Storage 只关心字节数,不关心字节怎么解释。dtype 信息全部下放到 TensorImpl::data_type_。这意味着同一段 StorageImpl 可以同时被一个 fp32 张量和一个 int32 张量持有(虽然语义上很危险),让 tensor.view(torch.int32) 这种”按位重解释”成为零拷贝操作。
如果你今天看 Python 文档会发现两个 API 共存:
tensor.storage()—— 老接口,返回 typed storage(其实底层是 untyped + 类型 wrapper),已弃用tensor.untyped_storage()—— v2.x 推荐的新接口,直接返回 untyped storage
新代码应该用后者。这是 PyTorch BC 政策的一个例子:老 API 仍能工作,但内部走 wrapper,新代码应该迁移。
2.5.1 为什么 Storage 也要分两层
你可能注意到一个奇怪的设计:Tensor 是 TensorBase + TensorImpl 两层,Storage 又是 Storage + StorageImpl 两层。为什么不直接合并?
答案是 API 稳定性。
Storage 与 Tensor 是用户可见的 C++ 公开 API,它们的 ABI(二进制接口)必须长期稳定 —— 不能因为内部加了一个字段就破坏所有依赖 libtorch 的 C++ 项目。但 StorageImpl / TensorImpl 是内部实现,可以自由演进。两层包装让”内部演进”和”外部稳定”得以共存。
这是 PyTorch 在所有公开类上都用的模式:
| 公开 API(稳定) | 内部实现(可演进) |
|---|---|
Tensor | TensorImpl |
Storage | StorageImpl |
Generator | GeneratorImpl |
RecordFunction | RecordFunctionImpl |
理解这个模式,你看 PyTorch 源码就不会迷惑”为什么每个东西都有 *Impl 后缀的兄弟”。
这种”包装类 + Impl”的二层结构在 C++ 圈子里有个名字叫 PIMPL(Pointer to IMPLementation)惯用法。优点是包装类的二进制表示永远是”一个指针 +(可能的)几个值类型字段”,编译期 ABI 极其稳定;缺点是每次访问字段都要走一次指针跳转,多一次潜在的 cache miss。PyTorch 选择 PIMPL 的代价是在张量层每次属性访问多 1-2 ns —— 但因为 dispatcher 自身的开销在百纳秒量级,这点 cost 完全淹没掉。
另一个隐藏的好处:Storage 这层薄包装让 PyTorch 能在 Python 端直接构造一个”指向已存在内存的张量”,而不需要 Python 看到内部的 StorageImpl。torch.from_blob(ptr, sizes, deleter) 这个 API 就是利用了这一点 —— 可以把外部库(OpenCV / cuDNN / 自家算子库)已分配的内存直接包装成 PyTorch 张量,零拷贝。这是 PyTorch 与 C++ 生态打通的关键 API 之一。
2.6 view 的实现:为什么不复制内存
回到本章开头的悖论。a.view(3, 4) 实际做的是:
sequenceDiagram
autonumber
participant U as 用户调用 a.view(3, 4)
participant TI_a as a 的 TensorImpl
participant TI_b as 新 TensorImpl
participant SI as 共享的 StorageImpl
U->>TI_a: a.view(3, 4)
TI_a->>TI_a: 检查 a 是否 contiguous
Note over TI_a: contiguous: 直接进入下一步<br/>非 contiguous: 抛 RuntimeError
TI_a->>TI_b: 新建 TensorImpl<br/>传入 storage_a (共享)
TI_b->>TI_b: sizes_ = (3, 4)<br/>strides_ = (4, 1)<br/>storage_offset_ = 0<br/>numel_ = 12
TI_b->>SI: storage_.intrusive_ptr<br/>refcount + 1
TI_b-->>U: 返回 b (Tensor 包装)
注意第 3 步:新 TensorImpl 的 sizes_ 和 strides_ 是重新计算的,不是从 a 复制的。strides 计算公式(行优先):
strides[i] = ∏(sizes[j] for j in i+1..ndim)
对 [3, 4]:strides = [4, 1](第 0 维步进 4 个元素,第 1 维步进 1 个元素)。
然后第 4 步:storage_ 的 intrusive_ptr 被复制(refcount +1)—— 数据本身没动一个字节。
这就是 view 的全部机制。一行代码:
// 简化的 view 实现,对应 aten/src/ATen/native/TensorShape.cpp 中的 view
auto new_impl = c10::make_intrusive<TensorImpl>(
Storage(storage_), // 共享同一个 storage
key_set_,
data_type_);
new_impl->set_sizes_and_strides(new_sizes, compute_strides(new_sizes));
new_impl->set_storage_offset(storage_offset_);
return Tensor(std::move(new_impl));
这种”共享 storage、独立元信息”的设计,让 PyTorch 中的 view、slice、transpose、permute、squeeze、unsqueeze 等几十个操作全部 O(1) 完成、零内存分配。这是 PyTorch 在张量层做的最重要的工程优化之一。
源码里实现这种共享有一个关键工具叫 copy_tensor_metadata(c10/core/TensorImpl.cpp 中实现)。它接受一个源 TensorImpl 和一个目标 TensorImpl,把元数据(sizes / strides / storage_offset / dtype / device / 各种 contiguous 缓存位)整体拷过去,但故意不拷 storage_ —— storage 是否共享由调用者决定。所有”视图类”操作的 C++ 实现底层都是这个函数加 storage 的 intrusive_ptr 复制。如果你写自定义算子需要造一个”看起来是张量的视图”,这个函数就是你的入口。
view 还有一个细节:它要求张量是 contiguous。这个限制不是任意的 —— 让我们用反证法理解为什么。假设有个张量 a,它的 sizes=[2,3] 但 strides=[1,2](一个伪装的 transpose),物理内存上元素顺序是 [a00, a10, a01, a11, a02, a12]。如果允许 a.view(6) 不复制内存,那么新视图按 strides=[1] 读出的顺序就是 [a00, a10, a01, ...] —— 这不是任何用户期望的”展平”。唯一的修复办法是先把 a copy 到 contiguous 布局,再 view —— 这正是 reshape 在背后偷偷干的事。
2.7 strides 与 offset:张量的”指针计算配方”
理解了 view,我们就能理解 PyTorch 索引的真正机制。
考虑一个 [3, 4] 的 fp32 张量 t,它的 storage_offset_ = 0,strides_ = [4, 1]。访问 t[1, 2] 时,PyTorch 计算:
byte_offset = (1 × strides[0] + 2 × strides[1] + storage_offset_) × sizeof(fp32)
= (1 × 4 + 2 × 1 + 0) × 4
= 24 bytes
然后从 storage.data_ptr() + 24 读出一个 fp32 —— 这就是 t[1, 2] 的值。
flowchart LR
subgraph idx["t[1, 2] 的访问"]
I["索引 (1, 2)"]
S["strides = [4, 1]<br/>storage_offset = 0"]
F["公式: 1×4 + 2×1 + 0<br/>= 6 个元素 = 24 字节"]
I --> F
S --> F
end
F --> P["data_ptr + 24"]
P --> V["读 fp32 → 得到值"]
style F fill:#fef3c7,stroke:#f59e0b,stroke-width:2px
关键认识:strides 是按”元素数量”算的,最后乘 sizeof(dtype) 才得到字节数。
2.7.1 transpose 不复制内存
理解了 strides,看一下 transpose 的魔法:
a = torch.arange(12).view(3, 4)
# sizes = [3, 4], strides = [4, 1], storage_offset = 0
b = a.t() # transpose
# sizes = [4, 3], strides = [1, 4], storage_offset = 0
# 数据完全没动!
transpose 做的全部工作就是 交换 sizes 和 strides 的对应位置。b[1, 2] 现在算出 1 × 1 + 2 × 4 = 9,而 a[2, 1] 算出 2 × 4 + 1 × 1 = 9 —— 指向同一个字节! 这正是 transpose 的数学定义。
代价是:b 的 strides [1, 4] 不再是”行优先递减”,所以 b.is_contiguous() == False。b 是一个非连续视图。
2.7.2 slice 用 storage_offset
a = torch.arange(12)
# sizes = [12], strides = [1], offset = 0
b = a[3:7]
# sizes = [4], strides = [1], offset = 3
# 数据共享,offset = 3 跳过前 3 个元素
storage_offset_ 字段就是为了支持 slice —— 不复制数据,只把”起点”往后挪。
2.7.2.5 NCHW vs NHWC:strides 在内存布局优化里的应用
考虑一个图像张量 [batch=8, channels=3, height=224, width=224]。两种合法的物理排布:
graph LR
subgraph NCHW["NCHW 布局 (标准 contiguous)"]
N1["sizes = [8, 3, 224, 224]<br/>strides = [3·224·224, 224·224, 224, 1]<br/>= [150528, 50176, 224, 1]"]
N2["内存按 NCHW 顺序排:<br/>所有 R 通道连续 → 所有 G → 所有 B"]
end
subgraph NHWC["NHWC 布局 (channels-last)"]
H1["sizes = [8, 3, 224, 224]<br/>strides = [3·224·224, 1, 3·224, 3]<br/>= [150528, 1, 672, 3]"]
H2["内存按 NHWC 顺序排:<br/>每个像素的 RGB 三通道连续<br/>→ 下个像素 RGB → ..."]
end
style N1 fill:#dbeafe
style H1 fill:#fef3c7
注意两个张量的 sizes 完全相同,但 strides 不同 —— 物理内存布局完全不同。channels-last 让”同一像素的不同通道”在内存中相邻,对某些卷积 kernel(特别是 NVIDIA Tensor Core 的 NHWC kernel、Apple Neural Engine 的卷积单元)友好得多,性能可以提升 30-50%。
PyTorch 通过 tensor.to(memory_format=torch.channels_last) 来切换布局。从源码角度,这个 API 做的事就是 重新分配 storage + 用不同的 strides 写入。从此该张量的 is_channels_last_ 位域被置 1,后续算子会优先选择 NHWC 路径。
这种”形状不变、strides 变”的能力,是 PyTorch 张量模型相对于 TensorFlow 1.x 等”形状即布局”框架的优势 —— 用户不需要重写代码,只通过 strides 切换内存布局就能享受硬件优化。
2.7.3 broadcast 用 stride = 0
广播是 PyTorch strides 模型最优雅的应用之一。torch.zeros(5).expand(3, 5) 怎么实现?
a = torch.zeros(5)
# sizes = [5], strides = [1]
b = a.expand(3, 5)
# sizes = [3, 5], strides = [0, 1] ← 注意 stride[0] = 0
b[i, j] 算出 i × 0 + j × 1 = j —— 不管 i 怎么变,都映射到同一行 a!stride = 0 把”伪复制”做成了 O(1) 视图操作。第 1 章 §1.4 提到的”广播不复制内存”,背后就是这个机制。
注意 expand 和 repeat 的微妙区别:
expand用 stride=0 做零成本视图,不能 inplace 修改(多个逻辑位置指向同一物理位置,inplace 会导致语义混乱)repeat真的复制内存,可以自由 inplace 修改
错把 repeat 当 expand 用是新手常见显存爆炸的来源 —— a.repeat(1024, 1, 1) 把 1MB 张量变成 1GB;正确的零成本写法应该是 a.expand(1024, *a.shape[1:])。
在 PyTorch 的 aten::expand 实现里(aten/src/ATen/native/TensorShape.cpp),核心代码就是创建一个新 TensorImpl,对要扩张的维度把 stride 强行置 0,sizes 改为目标值。零循环、零 kernel 调用、O(1) 完成。
2.8 contiguous 与 reshape:什么时候必须复制
view 操作有一个限制:只能用在 contiguous 张量上。这是为什么:
a = torch.arange(12).view(3, 4) # contiguous
b = a.t() # 非 contiguous (strides=[1,4])
c = b.view(12) # 报错!RuntimeError
为什么 b.view(12) 不能做?因为 b 在物理内存里的元素排列是 [0, 4, 8, 1, 5, 9, 2, 6, 10, 3, 7, 11](按 [1, 4] strides 读出的顺序),而 view(12) 要求读出的是线性顺序。只通过修改 strides,无论怎么调都不能从 [1, 4] 变成”读出来正好是线性” —— 必须真的搬数据。
PyTorch 的 reshape 就是为这个场景准备的:
// 简化的 reshape 实现逻辑
Tensor reshape(const Tensor& self, IntArrayRef shape) {
if (self.is_contiguous()) {
return self.view(shape); // 零复制
} else {
return self.contiguous().view(shape); // 复制后再 view
}
}
reshape 的潜规则:它可能复制内存,也可能不复制。这种”行为依赖输入”的 API 在工程上常常被批评(“不可预测”),但 PyTorch 选择了这条路是因为:
- 95% 场景张量是 contiguous,reshape 都能零复制
- 5% 场景张量非 contiguous,自动 contiguous 比要用户手动加
.contiguous()友好
如果你写训练代码遇到性能瓶颈,用 profiler 看 aten::contiguous 的调用 —— 一次大张量 contiguous 是几十毫秒级的开销,藏在某个 reshape 里很容易被忽略。
一个真实的工程案例:HuggingFace Transformers 早期版本的 attention 实现里有一句 q.view(bsz, num_heads, seq_len, head_dim),在 q 已经是 transpose(1, 2) 的非 contiguous 视图时会报错。社区的修法是改成 q.reshape(...),这就让 reshape 在背后悄悄触发了一次 contiguous —— 在大 batch 训练中这一刀就是几个百分点的吞吐损失。后来 v4.30+ 的改造里把这种隐式 contiguous 全部移到提前的”准备阶段”,让训练的 hot path 上不再出现 reshape,吞吐提升数据明显。这种从 view/reshape 选择带来的性能差距,正是理解三件套的实战回报。
另一个相关的踩坑场景:tensor.contiguous() 在 contiguous 张量上是 no-op(直接返回 self),但在非 contiguous 张量上会启动一次 kernel 复制。如果你不确定一个张量的状态,调用 .contiguous() 是安全的,但要意识到它可能不是免费的。reshape 是同样的逻辑:能 view 就 view,不能就 contiguous + view。这种”行为依赖输入”的 API 在调试时容易让人困惑 —— 一个解法是在性能敏感代码里显式用 view 并要求张量已 contiguous,让错误尽早暴露。
2.9 is_contiguous_ 是缓存而非计算
最后一个有意思的小细节。看 TensorImpl 的字段:
// c10/core/TensorImpl.h:2950 附近的位域
bool is_contiguous_ : 1;
bool is_channels_last_ : 1;
bool is_channels_last_contiguous_ : 1;
bool is_channels_last_3d_ : 1;
bool is_channels_last_3d_contiguous_ : 1;
bool is_non_overlapping_and_dense_ : 1;
这些都是缓存字段。每次 tensor.is_contiguous() 不重新计算 —— 直接读位域。
为什么这么做?因为 is_contiguous 在 PyTorch 内部被频繁查询(每次进 dispatcher、每次进 TensorIterator 都查),绝不能每次都遍历 strides 重算。但这就要求每次修改 strides 时手动同步缓存:
// c10/core/TensorImpl.cpp 中 set_sizes_and_strides 调用的 refresh_contiguous
void TensorImpl::refresh_contiguous() {
is_contiguous_ = compute_contiguous();
is_channels_last_ = compute_channels_last_contiguous_2d();
is_channels_last_contiguous_ = ...;
...
}
如果你写自定义算子直接修改了 tensor.strides_ 而忘了调 refresh_contiguous —— PyTorch 后续基于 is_contiguous_ 的所有判断都会错,引发非常隐蔽的 bug。这是 PyTorch 团队反复警告自定义算子作者”不要直接改 strides”的原因。
refresh_contiguous 的实现也展示了 PyTorch 处理多种内存格式的方式:除了普通的 contiguous(NCHW),还有 channels-last(NHWC,第 4 个维度连续)和 channels-last-3d(NDHWC)。这些格式各有专门的位域缓存,因为不同硬件后端(NVIDIA Tensor Core、AMD CDNA、Apple ANE)对内存布局有不同偏好。第 9 章会讲到 nn.Conv2d(..., memory_format=torch.channels_last) 这种 API,背后就是在张量层面把数据按 NHWC 排布以喂给 cuDNN 的 NHWC kernel。
为什么这种缓存机制的代价没有更频繁地显式显现?因为 PyTorch 内部修改 strides 的入口被严格收敛到 set_sizes_and_strides 一类的方法里,每个入口都强制调用 refresh_* 链。用户级 API 几乎不可能”跳过”这个机制。直接操作 TensorImpl 字段属于 PRIVATE API,源码里有大段警告评论。理解这个边界,你就理解了 PyTorch 在”灵活性”和”内部一致性”之间画的那条线。
第 22 章自定义算子会再讨论这个坑。
2.9.5 autograd_meta_ 的”按需付费”设计
回到 TensorImpl 的字段表,注意 autograd_meta_ 这个字段:
// c10/core/TensorImpl.h:2879 附近
std::unique_ptr<c10::AutogradMetaInterface> autograd_meta_ = nullptr;
它默认是 nullptr。只有在张量真的需要梯度时才被分配。这种”按需付费”的设计是 PyTorch 性能的隐藏功臣 —— 因为 inference 场景下绝大多数张量都不需要梯度,省掉这次堆分配能在每秒上百万次张量创建里省下几兆字节的内存与几百微秒的分配时间。
AutogradMetaInterface 是一个抽象基类,真正的实现是 torch::autograd::AutogradMeta(在 torch/csrc/autograd/variable.cpp 中)。这种”接口在 c10、实现在 torch”的拆分是因为 c10 是 PyTorch 的最底层库,不能依赖 autograd 的具体实现。但 TensorImpl 又必须能放下 autograd 元信息 —— 解法就是在 c10 层定义抽象接口,在 torch 层注入具体实现。第 7 章 autograd 那章会拆这个机制。
源码里还藏着一个 NOTE:
// c10/core/TensorImpl.h:2878 附近的注释
// 1. autograd_meta_ == nullptr
// 2. autograd_meta_ is default constructed (semantically, same as (1))
// 3. autograd_meta_ has nontrivial information content
也就是说,autograd_meta_ 既可以是 nullptr,也可以是默认构造的对象。所有 autograd 路径上的代码都必须先 null check,再访问字段。这种”防御性编程”在性能敏感代码里很常见 —— 把”通常情况零开销 + 偶发情况慢路径”的优化思路贯彻到每个字段。
2.9.6 PyObjectSlot:Python 与 C++ 的双向桥
另一个不常被讨论但极其重要的字段是 pyobj_slot_:
impl::PyObjectSlot pyobj_slot_;
它定义在 c10/core/impl/PyObjectSlot.h。作用是 建立 C++ TensorImpl 与 Python torch.Tensor 实例之间的双向引用。
为什么要双向?因为 PyTorch 既要支持”C++ 创建的张量被 Python 用”(C++ → Python:用 THPVariable_Wrap 把 TensorImpl 包成 PyObject),也要支持”Python 创建的张量被 C++ 用”(Python → C++:拿 THPVariable.cdata 拿到底层 TensorImpl)。如果不存这个双向链接,每次 C++ → Python 都要新建一个 PyObject,老的 PyObject 上挂的属性、Python 子类、__torch_function__ 注册信息全都丢掉。
pyobj_slot_ 的实现还涉及 GIL 与原子操作的精细同步 —— 因为这个 slot 可能被多个 Python 线程同时查询(PyTorch 默认释放 GIL 跑 C++ kernel)。源码里这部分有大量注释解释 race condition 的处理。第 5 章和第 9 章 nn.Module 都会回到这个字段。
2.9.7 TensorImpl 的子类生态
TensorImpl 不是 PyTorch 的唯一张量后端。它有几个重要的兄弟:
| 子类 | 路径 | 用途 |
|---|---|---|
SparseTensorImpl | aten/src/ATen/SparseTensorImpl.h | 稀疏 COO 张量,存 indices + values 而非密集数据 |
SparseCsrTensorImpl | aten/src/ATen/SparseCsrTensorImpl.h | 稀疏 CSR 张量,存 crow_indices / col_indices / values |
NestedTensorImpl | aten/src/ATen/native/nested/NestedTensorImpl.h | 嵌套张量,每行长度可变(用于变长序列) |
OpaqueTensorImpl<T> | aten/src/ATen/OpaqueTensorImpl.h | 不透明张量,把第三方库(如 MKLDNN、QNNPACK)的私有数据格式包成张量 |
UndefinedTensorImpl | c10/core/UndefinedTensorImpl.h | 单例”空张量”,作为可选 Tensor 参数的默认值 |
每个子类都重写了 TensorImpl 的部分虚函数(如 sizes_custom() / strides_custom() / numel_custom()),用 sizes_strides_policy_ 字段触发到自定义实现。这种”基类提供默认 + 子类按需 override”的扩展点是 PyTorch 支持非传统张量的关键机制。
特别提一下 UndefinedTensorImpl —— 它是一个全局单例(不是每次创建一个新对象)。所有”未定义”的 Tensor 参数(在某些算子里 grad、out 等可选参数没传时)都指向这个单例。这个 trick 让 PyTorch 处理”可选张量”参数时几乎零成本 —— 不用 std::optional<Tensor>,直接传 Tensor,调用方判断 is_defined() 即可。
2.10 一个真实的内存共享谱系
把这一章学到的所有东西串起来,看一段真实的训练代码:
import torch
x = torch.randn(8, 3, 224, 224, device='cuda') # NCHW input
# 原 storage: 8*3*224*224*4 = 4.8 MB
x_flat = x.view(8, -1) # [8, 150528]
# 共享 storage, 元数据复制
x_t = x_flat.t() # [150528, 8]
# 共享 storage, strides 翻转, 非 contiguous
x_first = x[0] # [3, 224, 224]
# 共享 storage, storage_offset = 3*224*224 = 150528
x_chw = x[0, :, :100, :100] # [3, 100, 100]
# 共享 storage, offset 不变,sizes/strides 改变
# 注意: 非连续 (因为最后两维只取了部分)
x_clone = x_first.clone() # [3, 224, 224]
# 新 storage! clone 是显式深拷贝
整段代码里,五个看起来不同的张量,有四个共享同一段 4.8 MB 显存,只有 x_clone 真的复制了一份。
这种”零成本视图”的设计是 PyTorch 训练效率的基石之一。Hugging Face Transformers 里大量的 view / transpose / reshape 操作之所以不会引爆显存,就因为它们绝大多数时候只是在改 TensorImpl 元信息,碰都不碰底层数据。
值得注意:这种”视图链”也带来一个常被忽略的内存现象 —— 只要还有一个视图存活,底层 storage 就不能被释放。考虑一段代码:
huge = torch.randn(10_000_000, device='cuda') # 40 MB
small = huge[:10] # 10 个元素的视图
del huge # huge 被删了,但...
# small 还在 → storage_impl 的 refcount 仍是 1 → 整 40MB 都被锁住
del huge 看似释放了大张量,但 small 通过共享的 StorageImpl 仍然让 40 MB 显存挂着。新手在写 dataloader / preprocessing 时容易踩这个坑:保留了一个小视图就把整大块内存锁住。正确做法是 small = huge[:10].clone(),让 small 走自己的 storage。
诊断这种”视图泄漏”的一个工具:torch.cuda.memory._record_memory_history() 加 torch.cuda.memory._dump_snapshot() 能产生一个时间序列的显存快照,一看就知道哪些 storage 因为视图而长期未释放。第 4 章和第 21 章会更深入这个工具。
PyTorch 团队也意识到这个问题,所以做了一个细节优化:tensor.detach() 创建的张量虽然共享 storage,但不共享 version_counter。这意味着 detach 后做 inplace 修改不会污染原张量的反向图。这是设计三件套时一个微妙的 trade-off —— 让 inplace 操作的”传染范围”和”内存共享范围”不完全重合。
graph LR
SI[StorageImpl<br/>4.8 MB on CUDA]
TI1["TensorImpl x<br/>sizes=(8,3,224,224)<br/>strides=(150528,50176,224,1)<br/>offset=0<br/>contiguous"]
TI2["TensorImpl x_flat<br/>sizes=(8,150528)<br/>strides=(150528,1)<br/>offset=0<br/>contiguous"]
TI3["TensorImpl x_t<br/>sizes=(150528,8)<br/>strides=(1,150528)<br/>offset=0<br/>NOT contiguous"]
TI4["TensorImpl x_first<br/>sizes=(3,224,224)<br/>strides=(50176,224,1)<br/>offset=150528<br/>contiguous"]
TI5["TensorImpl x_chw<br/>sizes=(3,100,100)<br/>strides=(50176,224,1)<br/>offset=150528<br/>NOT contiguous"]
TI6[TensorImpl x_clone<br/>独立 sizes/strides<br/>独立 offset]
SI2[新 StorageImpl<br/>独立 0.6 MB]
SI --> TI1
SI --> TI2
SI --> TI3
SI --> TI4
SI --> TI5
TI6 --> SI2
style SI fill:#fef3c7,stroke:#f59e0b,stroke-width:2px
style SI2 fill:#dcfce7,stroke:#22c55e
2.10.5 data_type_ 与 TypeMeta:dtype 在 C++ 里的表示
PyTorch 的 torch.float32、torch.bfloat16 这些 dtype 在 C++ 里对应的类型是 caffe2::TypeMeta,它定义在 c10/util/typeid.h 中。
TypeMeta 看起来只是一个 dtype 标识,但它内部还藏着关于这个类型的元信息:字节大小、构造函数指针、析构函数指针、复制函数指针。这是因为 PyTorch 早期允许”任意类型”的张量(包括 caffe2::Blob 风格的对象张量),需要为非 POD 类型保留类型擦除的构造 / 析构能力。今天 PyTorch 已经把支持的类型收敛到一组 POD 类型(fp32/fp16/bf16/int* 等),但 TypeMeta 这个类型擦除接口还保留着,作为历史遗产。
源码里 dtype 的注册集中在 c10/core/ScalarType.h,用一个宏批量定义所有合法 dtype:
// c10/core/ScalarType.h 中的 AT_FORALL_SCALAR_TYPES_AND7
#define AT_FORALL_SCALAR_TYPES_AND7(...) \
_(uint8_t, Byte) \
_(int8_t, Char) \
_(int16_t, Short) \
_(int, Int) \
_(int64_t, Long) \
_(at::Half, Half) \
_(float, Float) \
_(double, Double) \
_(at::BFloat16, BFloat16) \
...
这种”宏循环”在 PyTorch 里随处可见。第 5 章会展示 dispatcher 怎么用同一个宏在每个 dtype 上”展开生成 if-else 分发表” —— 这是 C++ 高性能多类型 dispatch 的标准技巧。
data_type_ 字段只占 2 字节,因为 TypeMeta 实际上只是一个 16-bit 的 ID(指向全局类型表的索引)。这个紧凑表示是 TensorImpl 字段压缩战的又一例子。
2.10.6 version_counter_:inplace 操作的”安全网”
另一个值得介绍的字段是 version_counter_:
c10::VariableVersion version_counter_;
这是 autograd 用来检测 inplace 修改的机制。每次张量发生 inplace 操作(如 x.add_(1)),它的 version_counter_ +1。autograd 在反向传播时如果发现某个保存的张量的 version 比保存时高,就抛错”one of the variables needed for gradient computation has been modified by an inplace operation”。
这是为什么 PyTorch 训练里那段经典报错来源。它的设计哲学是 “宁可早报错,不要悄悄出错” —— inplace 修改可能让反向传播得到错误梯度,但表面上训练还能继续,损失正常下降,模型却越训越差。version 检查把这种隐蔽错误暴露在第一时间。
源码上 c10::VariableVersion 是一个 intrusive_ptr<VersionCounter>,多个张量(如 view 链)共享同一个 counter —— 因为修改任意一个 view 都会污染整个共享 storage 的所有视图。第 7 章 autograd 会再回到这里。
2.11 横向对比:NumPy / JAX / TensorFlow 的张量模型
PyTorch 的”strides + offset + 共享 storage”模型并不是首创,它继承自 NumPy。但和其他 ML 框架对比,仍有有趣的差异:
| 框架 | 张量模型 | 视图机制 | 是否暴露 strides |
|---|---|---|---|
| NumPy | ndarray 单层 | 多视图共享 base | 是,且用户常直接操作 |
| PyTorch | Tensor → TensorImpl → Storage → StorageImpl 四层 | 多 TensorImpl 共享 Storage | 是,但用户少直接用 |
| JAX | 不可变 Array | 没有 view,每次操作要么是真的视图函数(lax.dynamic_slice)要么复制 | 否,strides 完全内部 |
| TensorFlow | tf.Tensor | 没有真正的内存共享 view | 否 |
| MindSpore | ms.Tensor | 类 PyTorch 视图 | 部分暴露 |
JAX 选择”完全不暴露 strides”是因为它走的是函数式不可变路线 —— 没有 inplace 操作就不需要担心”修改一个视图会污染另一个张量”。PyTorch 走的是”可变 + 视图”路线,必须暴露 strides 来支持高级用法(如自定义算子直接操作 tensor.stride())。
两条路线没有绝对优劣,但理解差异有助于你在多框架间切换时不被坑:JAX 用户来 PyTorch 第一次写 b = a.view(3, 4); b += 1,会被”咦 a 怎么也变了”震惊;PyTorch 用户去 JAX,会被”为什么 x.at[1].set(0) 返回新张量”震惊。
2.12 跨书关联:与丛书其他书的呼应
本章涉及的概念在丛书其他书里有相关讨论,可以对照阅读:
- 《vLLM 内核探秘》第 4-5 章 PagedAttention / KV Cache:vLLM 的 KV Cache 也大量使用”共享 storage + 不同视图”模式,每个 sequence 的 KV 是底层 block pool 的视图。理解了本章的 storage_offset / strides 机制,再看 vLLM 的 BlockTable 会顺畅很多
- 《Tokio 异步运行时》智能指针章:
intrusive_ptr与 Tokio 中的Arc都是 RAII 风格的引用计数,但 PyTorch 的intrusive_ptr把 refcount 嵌入对象本身而不是放在控制块里,理由会在第 3 章详细讲 - 《Rust 编译器之路》第 X 章 Layout:Rust 的
Layout也描述类似的”size + alignment”信息,但 Rust 是编译期已知,PyTorch 是运行期描述(因为张量形状到运行时才知道)
2.12.5 历史回顾:Tensor 和 Variable 的合并
值得回顾的一段历史:在 PyTorch 0.4 之前,Tensor 和 Variable 是两个不同的类:
torch.Tensor—— 不带梯度的多维数组torch.autograd.Variable—— 带梯度的张量,包了一层Tensor
这意味着用户必须手动 Variable(tensor) 才能让张量参与 autograd,写训练代码时充满了类型转换。源码层面也很丑 —— 每个算子要写两遍:一遍接 Tensor、一遍接 Variable。
PyTorch 0.4(2018 年 4 月)做了一次”统一” —— 把 Variable 合并到 Tensor,让每个张量天然能参与 autograd。这次合并在 TensorImpl 上留下的痕迹就是 autograd_meta_ 字段:原来 Variable 的”是否需要梯度 / 反向函数”等元数据,全部塞进 TensorImpl 的这个 std::unique_ptr 里。
理解这段历史,你就理解了为什么源码里仍然能看到 VariableType / VariableTypeManual.cpp 这类老命名 —— 它们是 0.4 合并的”考古遗迹”,PyTorch 维护者在缓慢清理但还没清完。读源码时遇到这些老名字,把它们当作”autograd 实现”理解就行,没必要被命名搞糊涂。
这次合并也是 PyTorch 走向”非侵入式 autograd”的关键 —— 用户可以用同一个 Tensor API 写代码,autograd 的开关由 requires_grad flag 与 torch.no_grad() 上下文控制,不需要切换到另一个类。这是 PyTorch 比 TensorFlow 1.x 用户体验更好的根本原因之一。
2.12.6 一段微基准:看视图操作有多便宜
理论说完,看一段实测:
import torch, time
n = 1_000_000
a = torch.randn(1024, 1024, device='cuda')
torch.cuda.synchronize()
t0 = time.perf_counter()
for _ in range(n):
b = a.view(-1) # 纯视图
torch.cuda.synchronize()
t1 = time.perf_counter()
t2 = time.perf_counter()
for _ in range(n):
b = a.clone() # 真复制
torch.cuda.synchronize()
t3 = time.perf_counter()
print(f'view per call: {(t1-t0)/n*1e6:.2f} us')
print(f'clone per call: {(t3-t2)/n*1e6:.2f} us')
在 H100 上的典型结果:
view per call: ~1.5 us # 只是创建 TensorImpl
clone per call: ~120 us # 4 MB H2H 拷贝 + 新 storage
差了两个数量级。这就是”视图廉价”的实证 —— 写训练代码时,能用 view 的地方绝不要用 clone。第 21 章会讲怎么用 profiler 自动发现”该用 view 却用了 clone”的代码。
值得提醒:本章讨论的”零成本视图”只是 CPU 端的元数据复制零成本,进入 dispatcher 走一遍流程仍然要 1-2 微秒 —— view 不是真的”免费”,只是相对于实际数据搬运便宜两个数量级。在小张量、超高频调用的代码里,视图操作的 dispatcher 开销也会成为瓶颈,这正是 torch.compile 试图消灭的东西。
2.13 一个练习:自己解开三件套
import torch, ctypes
a = torch.arange(12, dtype=torch.float32).reshape(3, 4)
print('a.shape =', a.shape)
print('a.stride() =', a.stride())
print('a.storage_offset() =', a.storage_offset())
print('a.untyped_storage().nbytes() =', a.untyped_storage().nbytes())
print('a.data_ptr() =', hex(a.data_ptr()))
b = a.t() # 非 contiguous
c = a[1:, 1:] # offset != 0
d = a.view(2, 6) # 共享 storage
# 验证 b、c、d 都和 a 共享同一个 storage
for x, name in [(b, 'b'), (c, 'c'), (d, 'd')]:
print(f'{name} shares storage: '
f'{x.untyped_storage().data_ptr() == a.untyped_storage().data_ptr()}')
# 直接读 storage 的字节
ptr = ctypes.cast(a.data_ptr(), ctypes.POINTER(ctypes.c_float))
print('Raw bytes as floats:', [ptr[i] for i in range(12)])
跑一遍,把每个张量的 (shape, stride, offset) 三元组写下来。如果你能预测出 b[2, 1] 在 storage 里的字节偏移、并且和直接读 a.data_ptr() + offset_bytes 的结果对得上 —— 你就把本章的所有内容内化了。
2.14 本章给后面章节的”接口”
读完本章你应该带着以下几个心智模型进入第 3 章:
Tensor是按值传递的轻包装 —— 看到Tensor add(Tensor a, Tensor b)不要担心拷贝开销,只是引用计数 +1- 元数据和数据彻底解耦 —— 任何只改 sizes/strides/offset 的操作几乎免费,任何需要新 storage 的操作都是潜在性能瓶颈
is_contiguous()是 O(1) 查询而非 O(n) 计算 —— 但这就要求 PyTorch 内部所有修改 strides 的入口都要刷新缓存autograd_meta_默认 nullptr —— inference 路径 vs training 路径的本质性能差异从这里开始- 每加一个 TensorImpl 字段都要反复掂量 —— 这是为什么 PyTorch 不轻易引入新元数据,rather 倾向于用
extra_meta_溢出区或新增 dispatch key
第 3 章会拆 c10 这个底层包:Device 的统一表示、DType 与 TypeMeta 的关系、DispatchKey 的设计、最重要的 —— intrusive_ptr 为什么不是 std::shared_ptr。这些抽象是 TensorImpl 字段背后的”零件供应商”,理解它们才算把张量这一层彻底打通。
评论 0
还没有评论,来说两句吧。
评论加载失败,刷新重试。