PyTorch ATen 低层张量库架构与引用计数“黄金法则”：基于 aten/src 的源码级解读

aten/src 是 PyTorch 中低层张量库（Tensor Libraries）与 ATen C++ 绑定的所在地，官方 README 给出了该目录的两大核心议题：源自 Torch 时代的低层库谱系，以及 C 库代码中必须手动维护的引用计数规则。读完本文，你将理解 ATen 与历史库（TH/THC 等）的关系、PyTorch 张量视图为何依赖引用计数，并掌握在 C/C++ 内核代码中正确管理指针生命周期的“黄金法则”，以及其在现代 c10::intrusive_ptr 实现中的落地细节。

aten/src：低层张量库的“根据地”

README 开篇即点明：

This directory contains the low-level tensor libraries for PyTorch, as well as the new ATen C++ bindings.

也就是说，aten/src 同时承载了两代东西：

1. ATen（All Tensors）：现代 PyTorch 的张量运算库，是 torch 顶层 API 之下的 C++ 实现主体；
2. 历史低层库的遗留物与命名遗产：一批可以追溯到最初 Torch 项目的库。

README 用一张简洁的表格交代了这些库的谱系：

缩写	全称	现状
TH	TorcH	历史库
THC	TorcH Cuda	历史库
THCS	TorcH Cuda Sparse	已废弃（defunct）
THNN	TorcH Neural Network	已废弃（defunct）
THS	TorcH Sparse	已废弃（defunct）

文档特别提示：“You'll also see these abbreviations show up in symbol names.” 这些缩写至今仍出现在符号名中。结合当前仓库可以印证这一点：

• aten/src/THC/ 目录下如今只剩 THCAtomics.cuh 与 THCDeviceUtils.cuh 两个头文件——THC 库本体已不复存在，仅以头文件形式留存，正对应谱系表中“历史库”的定位；
• 而真正的现代实现集中在 aten/src/ATen/ 下：core/ 提供 Tensor/IValue/Dispatcher 等基础设施，native/ 按后端（cpu、cuda、hip、mkl、vulkan、mps、quantized 等）组织具体算子实现，functorch/ 承载 vmap 等功能。

构建组织：从 CMake 看 ATen 的后端矩阵

aten/src/ATen/CMakeLists.txt 展示了 ATen 如何把多后端源码纳入构建：它通过 file(GLOB_RECURSE ...) 分别收集 core/、cuda/、hip/、mkl/、mps/、vulkan/ 等目录的头文件与实现文件，并用 configure_file 依据 USE_CUDA、USE_ROCM、USE_MAGMA 等开关生成 Config.h、CUDAConfig.h、HIPConfig.h 等配置头。可以推断，这正是 README 所说“多个变体（variants）”在现代构建系统中的体现：同一套 ATen 接口，在不同设备后端下链接不同的实现。

另外，aten/src/ATen/core/dispatch/README.md 说明 core/dispatch/ 存放的是 c10 dispatcher——“a single point through which we are planning to route all kernel calls”，即所有内核调用的统一入口。这与本文主题直接相关：当你在 Python 中调用 torch.add 时，请求最终经由这个调度器路由到 native/ 下某个后端的具体内核，而内核内部操作的内存生命周期，正由本文要讲的引用计数规则保护。

引用计数：视图机制逼出来的手动内存管理

PyTorch 的张量之所以能以不同视图共享同一块底层存储，靠的正是引用计数。README 中的例子非常经典：对 Tensor 调用 view() 时，会分配出一个维度不同的新张量，但它与原张量共享同一个 c10::StorageImpl：

// 示意：view() 之后，新旧张量共享同一 StorageImpl
// 原张量与视图都持有一份对该 Storage 的引用

这就意味着 C 库代码必须手动追踪引用计数。README 为此提炼出一条必须记住的规则：

Golden Rule of Reference Counting: You must either FREE or RETURN a pointer which was returned by a function whose name begins with new or which you called retain on. If you return this pointer, your function name must begin with new.

逐字翻译：

1. 凡是名字以 new 开头的函数返回的指针，或者你调用过 retain 的指针，你要么 FREE 它，要么 RETURN 它，二者必居其一；
2. 如果你选择把它返回出去，你的函数名必须以 new 开头——即把“所有权随指针转移”的约定写进函数命名里。

对于长函数，责任落在开发者身上：逐一检查每一处 new* 调用，确保每一个出口（exit point）都有与之配对的 free。这一约定本质上是用命名纪律（naming discipline）来约束 C 语言的裸指针生命周期，因为 C 层没有析构函数可以兜底。

例一：newIndices——new 前缀不意味着“新分配”

第一个例子是获取稀疏张量索引的引用。README 给出的用法：

THIndexTensor *indices = THSTensor_(newIndices)(state, sparse);
// ... do some stuff ...
THIndexTensor_(free)(state, indices);

两个要点：

• newIndices 以 new 开头，所以调用方必须在使用完毕（通常是函数末尾）时 free 它；
• 值得注意的是，newIndices 实际上并不新分配一个索引张量，它只是返回对已有索引张量的一个指针。即 new 前缀承诺的是“所有权转移”，而不是“内存分配”。
• README 同时警告：DO NOT directly access the member variables of the struct——不要直接访问结构体成员。

再看 newIndices 的实现：

THIndexTensor *THSTensor_(newIndices)(const THSTensor *self) {
  // ...
  return THIndexTensor_(newNarrow)(self->indices, 1, 0, self->nnz);
}

它内部调用了 newNarrow（同样是 new 前缀），但没有 free 其返回值，而是直接返回——这恰好为它自己名字里的 new 提供了正当性：引用计数沿调用链向上传递，最终由最外层的调用方负责 free。这是一个典型的“所有权沿 new 命名链上溯”的模式。

例二：setStorage——把对象交给别人，不等于卸责

一个常见的直觉误区是：把指针传给了另一个函数，自己就不用管释放了。README 明确否定了这一点——如果对方函数需要长期持有该指针，它会自己 retain：

THByteStorage *inferred_size = THByteStorage_newInferSize(size, numel);
THTensor_(setStorage)(self, tensor->storage, tensor->storageOffset, inferred_size, NULL);
c10::raw::intrusive_ptr::decref(inferred_size);

模式拆解：

1. newInferSize 返回一个新引用（refcount = 1）；
2. setStorage 内部会自行 retain 一份（此时 refcount = 2），调用方手里仍持有第 1 份；
3. 因此函数末尾用 c10::raw::intrusive_ptr::decref(inferred_size) 放掉自己的那份（refcount 回到 1，由 Storage 持有）。

这里出现了现代机制 c10::raw::intrusive_ptr::decref 的身影——历史上的 free/retain 宏在当代代码库中已经由 intrusive_ptr 的 incref/decref 接管。

例三：newContiguous / retain——“统一释放点”技巧

第三个例子更微妙：你手里已有一个张量，但只有在某些条件下才需要调用 newContiguous 把它转成正确形态：

if (!(k_->stride(3) == 1) || !(k_->stride[2] == k_->size(3))) {
  kernel = THTensor_(newContiguous)(k_);
} else {
  THTensor_(retain)(k_);
  kernel = k_;
}
...
c10::raw::intrusive_ptr::decref(kernel);

两个分支分别拿回一个“属于自己的新引用”：条件分支来自 newContiguous 的返回，否则分支则显式 retain 一份。这样函数末尾就可以无条件 decref(kernel)，不必再写 if/else 收尾。README 点出了这一技巧的直觉：你希望条件表达式可以被一个等价的函数调用（如 kernel = THTensor_(newContiguous2D)(k_)）替换而不破坏生命周期正确性——即“拿到什么形式的引用，就用统一的方式还回去”。

README 的三条实战 Tips

原文档最后给出了三条经验，这里完整继承并展开：

Tip 1：提前 return 之前先释放。 如果函数中有早退（early exit），别忘了 free 到目前为止分配的所有指针。更好的做法是把早退判断移到分配之前，从而根本不需要清理代码。这与现代 C++ 中“尽早 fail-fast、推迟资源获取”的原则一脉相承。

Tip 2：move 是唯一的例外。 极少数情况下，“摧毁”输入能让算法更高效——这就是 move：把对象 move 走之后，不得再 free 它。这是黄金法则的唯一例外，且 README 强调当时整个代码库中只有一处 move 的实例。

Tip 3：THError 与错误路径。 库用 THError 报告错误。默认情况下它直接 abort 整个进程，所以调用 THError 之前不必确保所有指针都已释放；但良好的风格是在任何分配发生之前就调用 THError。原因是：某些路径（尤其是测试套件）会抛出 C++ 异常并尝试恢复，此时未释放的资源才会真正泄漏。

现代实现印证：c10::intrusive_ptr 与组合引用计数

README 中出现的 c10::raw::intrusive_ptr::decref 在 c10/util/intrusive_ptr.h 中有完整实现，这也是当前 PyTorch 张量对象（at::Tensor、Storage、TensorImpl 等）生命周期管理的实际机制。几个与上文规则直接对应的实现细节：

• 组合引用计数位布局：文件开头定义了位打包常量，uint64_t 中低 32 位是强引用计数（refcount），高 32 位是弱引用计数（weakcount），第 63 位标记对象是否带有 PyObject 包装（kHasPyObject）。refcount()、weakcount()、has_pyobject() 等辅助函数（c10/util/intrusive_ptr.h#L60-L74）负责解包。
• is_uniquely_owned：判断“引用计数恰好为 1 且无弱引用”的快速路径，用于避免不必要的拷贝等优化，正是黄金法则在运行时的高效化——当且仅当唯一拥有时，才能安全地“move 式”地复用缓冲。
• 原子 incref/decref：raw::intrusive_ptr::incref / decref（c10/util/intrusive_ptr.h#L1159-L1176）内部使用 std::atomic<uint64_t> 的组合计数做无锁增减，decref 到零时析构对象。这解释了为什么 C++ 侧的张量传递可以安全地在多线程间流转，而 C 侧遗留代码仍需按黄金法则手动配对 free。
• PyObject 联动：实现中当引用计数从 1 升到 2 时会对底层 PyObject 做 incref，从 2 降到 1 时 decref（见文件内注释与 incref_pyobject/decref_pyobject 虚函数），保证 C++ 张量与 Python 对象同生共死。

换言之，README 描述的那套 new/free/retain 命名纪律是 C 库时代的“人工内存安全”，而今天的 c10::intrusive_ptr 把同样的语义（“持有即加引用，交出引用即可能析构”）固化进了类型系统——你不需要再靠函数名前缀来猜测所有权，但理解这套历史约定，依然是读懂 ATen 遗留代码与命名规范（大量 new* 风格 API 名）的钥匙。

小结

aten/src 的 README 篇幅不长，信息密度却很高：它交代了 ATen 之下 TH/THC/THCS/THNN/THS 的库谱系与废弃状态，给出了一张从 aten/src/THC/ 仅存两个头文件、到 aten/src/ATen/ 庞大实现矩阵的演化地图；更重要的是，它用“黄金法则 + 三个代码样例 + 三条 Tips”完整定义了在共享存储（c10::StorageImpl）的视图机制下手动管理引用计数的方法论。结合 c10/util/intrusive_ptr.h 的组合引用计数实现与 core/dispatch 的调度器设计，这条从“裸指针纪律”到“侵入式智能指针”的演化线，正是理解 PyTorch C++ 内核内存管理的关键脉络。