PyTorch Tensor 与 TensorImpl 机制解读
1. 为什么要把 Tensor 和 TensorImpl 分成两层
初次接触 PyTorch 源码的人常有一个疑惑：Python 里 torch.Tensor 明明是一个对象，为什么 C++ 里又有 at::Tensor 和 c10::TensorImpl 两个类？
答案是值语义 vs 数据本体的分离：
at::Tensor：一个轻量的"引用/handle"，本身几乎不存储数据，可以随意拷贝、传参、放进容器，开销只是一次引用计数的加减。c10::TensorImpl：真正持有 sizes/strides/dtype/device/storage 等元信息的对象，在堆上只有一份。
多个 at::Tensor 可以指向同一个 TensorImpl（比如 view 操作、赋值拷贝），它们共享同一份元信息和底层数据，改一个，另一个也能看到变化——这正是 PyTorch “view 语义"和"就地修改语义"的根基。
at::Tensor a; at::Tensor b = a; // 只拷贝了指针，refcount+1
│ │
└──────────┬─────────────┘
▼
c10::TensorImpl (堆上唯一一份)
├── sizes/strides
├── storage_ ──────────► StorageImpl（真正的数据 buffer）
├── dtype / device
└── ...
这种设计在 C++ 里通过 c10::intrusive_ptr<TensorImpl> 实现：at::Tensor 内部只有一个 intrusive_ptr<TensorImpl> 成员，拷贝 Tensor 就是拷贝这个智能指针。
2. intrusive_ptr：为什么不用 std::shared_ptr
PyTorch 没有直接用 std::shared_ptr<TensorImpl>，而是自己实现了 c10::intrusive_ptr。核心原因是性能：
std::shared_ptr需要额外分配一个控制块（refcount），对象和控制块是两次内存分配，且引用计数的加减要过一层间接寻址。intrusive_ptr要求目标类型（TensorImpl）自己继承c10::intrusive_ptr_target，把 refcount 直接嵌入对象本身。这样创建TensorImpl时只需一次分配，Tensor的拷贝/析构只是对这个内嵌计数器做原子加减，没有额外的控制块开销。
Tensor 在 PyTorch 里是被极高频次拷贝、传参的对象（几乎每个 op 调用都会拷贝若干个 Tensor 参数），这个优化在整体性能上是必要的。
3. TensorImpl 里到底存了什么
抛开各种为兼容历史遗留和特殊场景（quantized/sparse/nested 等）而存在的字段，TensorImpl 的核心成员大致是：
// c10/core/TensorImpl.h（节选，做了简化）
class C10_API TensorImpl : public c10::intrusive_ptr_target {
protected:
Storage storage_; // 指向底层数据 buffer
private:
std::unique_ptr<AutogradMetaInterface> autograd_meta_; // autograd 相关（grad/grad_fn等），默认懒分配
std::unique_ptr<ExtraMeta> extra_meta_; // named tensor 等不常用元信息，同样懒分配
protected:
c10::VariableVersion version_counter_; // 版本号，用于 autograd 的就地修改检测
impl::PyObjectSlot pyobj_slot_; // 与 Python 端对象的关联（PyObject* 缓存）
c10::impl::SizesAndStrides sizes_and_strides_; // shape 与 stride
int64_t storage_offset_ = 0; // 相对 storage 起始地址的偏移（以元素计）
int64_t numel_ = 1; // 元素总数
caffe2::TypeMeta data_type_; // dtype 元信息
std::optional<c10::Device> device_opt_; // 设备信息
DispatchKeySet key_set_; // 决定这个 tensor 该走哪些 dispatch 分支
...
};
几个关键点：
storage_是分离的一层：TensorImpl不直接管内存，而是持有一个Storage（对StorageImpl的引用计数包装）。这就是为什么x.view(...)、x[1:3]、x.t()这些 view 操作可以做到"零拷贝”——它们创建一个新的TensorImpl，改一下sizes_and_strides_和storage_offset_，但storage_指向同一份底层数据。sizes_and_strides_是 shape/stride 的核心存储，这也是 tensor 支持任意维度、任意 stride（包括负 stride、broadcast 用的 0 stride）的地方。autograd_meta_/extra_meta_都是懒分配的unique_ptr：绝大多数 tensor 不需要 autograd 或额外元信息，用指针而不是内嵌成员可以让"不需要这些特性的 tensor"不为它们付出内存代价。key_set_（DispatchKeySet）是分发机制的关键：每个 tensor 携带一组 dispatch key（如CPU、CUDA、Autograd、Sparse等），调用一个 op 时，PyTorch 的 dispatcher 根据参与运算的所有 tensor 的key_set_的并集，决定这次调用具体落到哪个 backend 的 kernel 实现、要不要先经过 autograd 层记录计算图。这也是 PyTorch 能同时支持 CPU/CUDA/autograd/量化等多套机制正交组合的基础。
4. Storage / StorageImpl：数据本体又是另一层
TensorImpl 里的 storage_ 是 c10::Storage，它是对 c10::StorageImpl 的 intrusive_ptr 包装，StorageImpl 才是真正持有裸内存的对象：
// c10/core/StorageImpl.h（节选）
struct C10_API StorageImpl : public c10::intrusive_ptr_target {
at::DataPtr data_ptr_; // 真正的裸内存指针（带自定义 deleter）
SymInt size_bytes_; // 字节数
bool resizable_;
at::Allocator* allocator_; // 分配器（不同 device 有不同实现）
...
};
为什么又要单独抽一层 Storage？因为"元信息"（shape/stride/dtype）和"数据本体"的生命周期、共享方式并不总是一致：
- 一个
Tensor做view/reshape/切片后，会产生新的TensorImpl（不同的 shape/stride/offset），但共享同一个StorageImpl（同一份底层数据）。 Storage这一层还承担了 CPU tensor 与 numpy 数组共享内存、torch.from_blob包装外部内存等场景，这些场景下"元信息"完全是 PyTorch 自己的，但"数据"来自外部，两层分离让这类需求可以干净地表达。
5. 自定义 TensorImpl：为什么以及怎么做
PyTorch 允许第三方后端（backend）继承 c10::TensorImpl 写自己的子类，这是"外接后端"（out-of-tree backend）机制的核心之一。常见动机是：在标准元信息之外，给 tensor 附加一些只有这个后端才关心的状态，同时复用 PyTorch 整套 dispatcher/autograd/view 机制。
一个最小的自定义 TensorImpl 大致要做这几件事：
class MyTensorImpl : public c10::TensorImpl {
public:
MyTensorImpl(c10::Storage&& storage, const caffe2::TypeMeta& dtype)
: c10::TensorImpl(std::move(storage), key_set_for_my_backend(), dtype) {}
// 1. shallow_copy_and_detach：detach() 等操作要用到，
// 要求返回"同类型"的一个新 TensorImpl，共享 storage，但独立的 autograd 身份。
c10::intrusive_ptr<c10::TensorImpl> shallow_copy_and_detach(
const c10::VariableVersion& version_counter,
bool allow_tensor_metadata_change) const override;
// 2. shallow_copy_from：set_data() 等操作会调用，
// 把另一个 impl 的元信息浅拷贝到当前对象（storage 指针、sizes/strides 等），
// 但保留当前对象的 autograd 身份（version_counter 由调用方决定是否替换）。
void shallow_copy_from(const c10::intrusive_ptr<TensorImpl>& impl) override;
private:
// 这里可以放任意自定义状态：缓存标记、后端私有 handle、统计信息……
MyCustomState state_;
};
shallow_copy_and_detach 和 shallow_copy_from 是必须重写的两个虚函数（PyTorch 源码里专门有一段 NOTE [ TensorImpl Shallow-Copying ] 注释解释语义），原因很直接：PyTorch 内部（尤其 autograd 引擎和 Tensor.detach() / Tensor.set_data()）需要在"不知道具体子类是什么"的情况下，也能正确地复制出一个同类型的新实例，而不是意外退化成基类 TensorImpl。如果不重写，子类特有的状态在 detach/set_data 之后就会丢失或类型错乱。
除此之外，创建自定义 TensorImpl 实例要用 at::detail::make_tensor<T>(...) 而不是直接 new：
at::Tensor t = at::detail::make_tensor<MyTensorImpl>(std::move(storage), dtype);
这个工厂函数负责把 T 的实例包进 intrusive_ptr，再包成 at::Tensor，保证引用计数和析构路径正确。
5.1 展开看这一行到底做了什么
make_tensor<T> 不是什么特殊魔法，顺着调用链往下看，其实是三层很朴素的包装（以 TensorBase 为例，Tensor 的逻辑一致）：
// aten/src/ATen/core/TensorBase.h
namespace detail {
template <typename T, typename... Args>
TensorBase make_tensor_base(Args&&... args) {
return TensorBase(c10::make_intrusive<T>(std::forward<Args>(args)...));
}
}
再往下，make_intrusive<T> 只是对 intrusive_ptr<T>::make 的转发：
// c10/util/intrusive_ptr.h
template <class TTarget, class... Args>
static intrusive_ptr<TTarget> make(Args&&... args) {
return intrusive_ptr(new TTarget(std::forward<Args>(args)...));
}
也就是说最终确实是 new T(...)，并没有绕开 new——真正的关键在**new 出来之后立刻发生了什么**。看 intrusive_ptr 接管这个裸指针时的构造函数：
explicit intrusive_ptr(TTarget* target)
: intrusive_ptr(target, raw::DontIncreaseRefcount{}) {
if (target_ != NullType::singleton()) {
TORCH_INTERNAL_ASSERT_DEBUG_ONLY(
target_->refcount_ == 0 && target_->weakcount_ == 0,
"...");
target_->refcount_.store(1, std::memory_order_relaxed);
target_->weakcount_.store(1, std::memory_order_relaxed);
}
}
这里有一个隐含的强约束：新 new 出来的对象，refcount_/weakcount_ 必须是 0（因为 TensorImpl 继承的 intrusive_ptr_target 里这两个计数器默认初始化为 0），构造函数看到之后把它们手动置为 1，代表"现在有且仅有这一个 intrusive_ptr 持有它"。这一步如果不做，或者自己 new 之后随手包一个 intrusive_ptr，计数器初值不对，后续的引用计数加减就全错了——这也是为什么源码专门写了断言防止在自定义构造函数里"意外"提前创建了指向 this 的 intrusive_ptr（注释里点名了这种误用）。
最后一层，TensorBase(intrusive_ptr<TensorImpl, ...> tensor_impl) 构造函数只是把这个已经计数正确的智能指针存进 impl_ 成员，并检查非空：
explicit TensorBase(c10::intrusive_ptr<TensorImpl, UndefinedTensorImpl> tensor_impl)
: impl_(std::move(tensor_impl)) {
if (impl_.get() == nullptr) {
throw std::runtime_error("TensorImpl with nullptr is not supported");
}
}
串起来就是：
make_tensor<T>(args...)
└─ make_tensor_base<T>(args...)
└─ TensorBase( make_intrusive<T>(args...) )
└─ intrusive_ptr<T>::make(args...)
└─ intrusive_ptr( new T(args...) ) // 唯一一次 new
└─ refcount_ = 1; weakcount_ = 1;
一次 new 分配对象本体，紧接着由 intrusive_ptr 的构造函数把内嵌的计数器初始化为 1，再把这个已经"计数正确"的指针搬进 TensorBase/Tensor 的 impl_ 字段。整个过程只有一次堆分配（对象和计数器是同一块内存，因为计数器就是 intrusive_ptr_target 的成员），这也是前面第 2 节提到 intrusive_ptr 比 shared_ptr 省一次控制块分配的具体落地位置。
6. 自定义 TensorImpl 能挂载什么样的状态
只要不破坏上面两个虚函数的语义，子类里几乎可以放任何自定义成员：布尔标记、缓存指针、指向额外 side-table 的 handle、状态机对象等等。工程上有两个常见考量：
- ABI 稳定性：如果这个 TensorImpl 子类会被编译进独立的动态库、又被另一个独立编译的模块使用，直接暴露成员变量会导致"新增/删除成员"就破坏二进制兼容。常见解法是 PIMPL（Pointer to IMPLementation）——头文件里只暴露一个
unique_ptr<Impl>和一组访问函数，真正的成员藏在.cpp里的Impl结构体中，日后增删这些私有成员不影响已编译好的调用方。 - 状态跟 tensor 生命周期绑定：把状态放进 TensorImpl 成员，天然享有和 tensor 完全一致的生命周期——tensor 被析构，状态自动跟着释放，不需要额外维护一个"tensor → 状态"的全局哈希表、也不需要担心忘记手动清理导致内存泄漏。这是相比"用外部 map 关联 tensor 和状态"更稳健的一种模式。
7. 小结
at::Tensor是值语义的引用句柄，真正的元信息全部在堆上的c10::TensorImpl里，两者靠c10::intrusive_ptr关联，拷贝Tensor只是引用计数的加减。TensorImpl本身也不直接持有数据，数据在Storage/StorageImpl里，这一层分离是 view/reshape 等零拷贝操作和 tensor 与外部内存互操作的基础。key_set_（DispatchKeySet）是 tensor 参与算子分发的"路由标签"，决定一次调用具体落到哪个 backend kernel、是否经过 autograd。- 自定义后端可以继承
TensorImpl附加私有状态，代价是必须正确实现shallow_copy_and_detach/shallow_copy_from，并通过at::detail::make_tensor<T>来创建实例；把状态挂在 TensorImpl 上的核心好处是生命周期自动跟随 tensor，不需要额外的全局注册表。