跳到主要内容

数据截至 (上游 commit c187ef3271d5)

01 · Tensor 与 Storage：一块内存的两层视图

这一章讲什么： 你手里的 torch.Tensor 到底是哪些 C++ 对象的组合。读完你会理解：为什么 t.view(3, 4)、t[2:]、t.t() 都不复制数据，为什么 x is y 为 False 的两个 Tensor 却共享同一块显存。

1. 它要解决的小问题​

深度学习要的「张量」和一段裸内存之间隔着一大堆元数据：形状、每维步长、数据类型、在哪台设备上、要不要梯度。

如果每次转置、切片都复制一份数据，显存带宽会被拖垮。所以 PyTorch 把「数据」和「怎么读数据」彻底分开：数据只有一份，元数据想造多少份都行。

2. 三层结构：Python 壳 → TensorImpl → Storage​

一个 torch.Tensor 从上到下是三层：

Python 层    torch.Tensor            你在 .py 里 import 的那个类
                │ (继承)
             torch._C.TensorBase     CPython 对象 THPVariable（带 ob_base 头）
                │ (内部持有)
C++ 层       TensorImpl             元数据：sizes/strides/dtype/device/key_set
                │ (intrusive_ptr 引用)
             StorageImpl            数据：data_ptr_ + size_bytes_ + allocator_

三层的对应代码：

层类型位置
Python 类class Tensor(torch._C.TensorBase)torch/_tensor.py:102
CPython 包装struct THPVariabletorch/csrc/autograd/python_variable.h:17
元数据对象struct TensorImpl : public c10::intrusive_ptr_targetc10/core/TensorImpl.h:510
数据缓冲struct StorageImpl : public c10::intrusive_ptr_targetc10/core/StorageImpl.h:55

两个细节值得记住：

  • 引用计数用的是 intrusive_ptr（计数器存在对象体内，而不是像 shared_ptr 那样外挂控制块），所以 Tensor 的拷贝/传递非常便宜。
  • Python 壳和 C++ 对象可以互相「复活」：TensorImpl 上有一个 pyobj_slot_（c10/core/TensorImpl.h:2921）记着它对应的 Python 对象，C++ 侧拿到裸 TensorImpl 时能找回同一个 Python 对象。

3. TensorImpl 的字段：一张「怎么读内存」的说明书​

TensorImpl 的成员集中在 c10/core/TensorImpl.h:2888 附近，挑关键的看：

字段行号干什么
Storage storage_c10/core/TensorImpl.h:2888指向数据缓冲（多个 Tensor 可共享同一个）
SizesAndStrides sizes_and_strides_c10/core/TensorImpl.h:2923每维大小 + 步长
int64_t storage_offset_c10/core/TensorImpl.h:2925本 Tensor 从缓冲的第几个元素开始读
int64_t numel_c10/core/TensorImpl.h:2930元素总数（缓存值）
TypeMeta data_type_c10/core/TensorImpl.h:2934元素类型（float32 等）
std::optional<Device> device_opt_c10/core/TensorImpl.h:2948所在设备
DispatchKeySet key_set_c10/core/TensorImpl.h:3045决定 dispatcher 路由的位掩码（第 2 章主角）
unique_ptr<AutogradMetaInterface> autograd_meta_c10/core/TensorImpl.h:2908梯度相关元数据（第 3 章主角）

地址计算公式就是这套元数据的全部语义：

element(i0, i1, ...) 在内存中的位置
  = data_ptr + (storage_offset + Σ ik * stridek) * sizeof(dtype)

源码注释也这么解释 storage_offset：大多数 Tensor 是 0，「对 Tensor 做索引切片的那个」会有非零 offset（c10/core/TensorImpl.h:744-745）。

4. 直觉演示：view 为什么免费​

用一段示意代码把上面三层的关系演出来（示意，非源码）：

# 示意，非源码
a = torch.arange(12)        # 分配一块 12 元素的 Storage
b = a.view(3, 4)            # 新建 TensorImpl：同一个 Storage，
                            #   sizes=[3,4], strides=[4,1], offset=0
c = a[2:]                   # 新建 TensorImpl：同一个 Storage，
                            #   sizes=[10], strides=[1], offset=2
b[1, 1] = 99                # 改的是共享内存 → a[5] 也变成 99

重点看：view 和切片只新建了一个 TensorImpl（几百字节的元数据），Storage 引用计数加一，数据一个字节都没动。

真实实现里，as_strided 系列就是干这个的：as_strided_tensorimpl（aten/src/ATen/native/TensorShape.cpp:1328）按你给的 size/stride/offset 造一个新的 TensorImpl。

5. 两个藏得很深的性能小心思​

5.1 小形状数组内联在对象体内​

绝大多数张量 ≤ 4~5 维。如果 sizes/strides 每次都去堆上 malloc，高频算子会被分配器拖慢。

SizesAndStrides 的做法是小尺寸优化：

// c10/core/impl/SizesAndStrides.h:10
#define C10_SIZES_AND_STRIDES_MAX_INLINE_SIZE 5

维度 ≤5 时，sizes 和 strides 各 5 个 int64 直接放在对象内的 inlineStorage_（c10/core/impl/SizesAndStrides.h:327）；超过才退化为堆分配（isInline() 判断在 c10/core/impl/SizesAndStrides.h:289）。

5.2 不要梯度的 Tensor 连 autograd 元数据都不分配​

autograd_meta_ 是个 unique_ptr，nullptr 就等价于「默认构造、不需要梯度」。源码注释明确说这是优化（c10/core/TensorImpl.h:2890-2908）：

autograd_meta_ can be nullptr, as an optimization. When this occurs, it is equivalent to having an autograd_meta_ pointing to a default constructed AutogradMeta.

推理场景下海量 Tensor 都是 requires_grad=False，这一手省掉了每个 Tensor 一次堆分配。

6. Storage：字节缓冲 + 分配器​

StorageImpl 的构造一目了然（c10/core/StorageImpl.h:55-70）：它只装三样东西——

字段干什么
data_ptr_指向原始字节（DataPtr 自带 deleter，跨设备/自定义分配器统一抽象）
size_bytes_缓冲有多少字节（注意：是字节数，不是元素数）
allocator_谁分配的（CPU/CUDA/pinned memory 各有一套），可 resize 时必须非空

由此可以回答一个常见困惑：「两个 Tensor 共享 storage」到底是什么意思？ 就是它们的 TensorImpl::storage_ 指向同一个 StorageImpl——共享的单元是「一整块字节缓冲」，而各自的 sizes/strides/offset 决定各自看见它的哪一段、以什么形状看。

7. 坑与边界​

  • view 不是总有定义的。 transpose 后内存不再连续，view 会直接报错——因为它只能表达「strides 可描述」的形状变化，任意重排得用 reshape（必要时复制）。这是元数据模型的代价。
  • 共享 storage 意味着互相看见修改。 对切片做 in-place 写会改到原 Tensor；autograd 为此配了 version counter 检测冲突（第 3 章）。
  • numel_ 是缓存的标量（c10/core/TensorImpl.h:2930），符号形状（SymInt）场景下另有 sym_numel 路径，别假设所有形状信息永远是编译期可知的小整数。
  • 看不出来：具体每个分配器（如 CUDACachingAllocator 的分块策略）怎么复用显存——那是 allocator 子系统，不在本章主线。

8. 本章代码地图​

主题文件路径符号名
Python Tensor 类torch/_tensor.pyTensor
CPython 包装对象torch/csrc/autograd/python_variable.hTHPVariable
元数据结构c10/core/TensorImpl.hTensorImpl、storage_、sizes_and_strides_、storage_offset_
小尺寸优化c10/core/impl/SizesAndStrides.hC10_SIZES_AND_STRIDES_MAX_INLINE_SIZE、isInline
数据缓冲c10/core/StorageImpl.hStorageImpl、data_ptr_、size_bytes_
view 的构造aten/src/ATen/native/TensorShape.cppas_strided_tensorimpl