TensorPlay
The three-layer memory stack: bytes, containers, typed tensors

A tensor, simply put, is “a piece of memory, plus attributes like shape, strides, and device”.

The core idea of the tensor system is: the underlying bytes must be shared, but the views must be independent. b = a[0] can be the same piece of memory with a different shape; conversely, a + 1 is new memory. View and clone — it’s hard for a single “tensor class” to do both well.

TensorPlay designs a three-layer representation model.

Why three layers

Suppose there were only one layer: the Tensor class directly holds a pointer obtained from malloc, plus shape, strides, and device. It’s hard to handle the following cases:

  • b = a[0] shares bytes, but has a different shape.
  • On CPU the memory comes from the caching allocator, on CUDA from cudaMalloc, and there are also occasions where it wraps a foreign buffer. The release logic is heterogeneous.
  • The gradient history of automatic differentiation is dropped on copy. One tensor class becomes complicated to manage.

The three layers: Storage -> TensorImpl -> Tensor

Start with the bottom-most piece, DataPtr:

class DataPtr {                        // move-only 句柄
  void*        data_;                  // 交给调用者的数据指针
  void*        ctx_;                   // 传给删除器的上下文(未必等于 data_)
  DeleterFnPtr deleter_;               // void(*)(void*):裸函数指针,无捕获
  Device       device_;                // 设备随分配一起走

  DataPtr(const DataPtr&) = delete;            // 拷贝被显式删除
  DataPtr& operator=(const DataPtr&) = delete;
  DataPtr(DataPtr&& other) noexcept : DataPtr() { swap(other); }
  ~DataPtr() { clear(); }
  void clear() { if (deleter_) deleter_(ctx_);
                 data_ = ctx_ = nullptr; deleter_ = nullptr; }
  void* release() { void* p = data_;          // 调用者接管所有权
                    data_ = ctx_ = nullptr; deleter_ = nullptr; return p; }
};

The deleter: for the same stretch of bytes, what came from malloc needs free, what came from new[] needs delete[], what came from cudaMalloc needs to go back through device memory’s release interface, and memory handed in from outside must not be touched at all. The destruction method is tied to the allocation, so it becomes deleter_, packed into the same handle, moving along with data_ and disappearing along with it; ctx_ is the argument passed to it at destruction time — most of the time that’s the data pointer itself, passed separately only when extra state has to come along.

Copying: the deleter carries ownership, so copying it once is a double free once. If you want to transfer, call release(); if you’re wrapping foreign memory, hang a deleter that does nothing. Tensors can be shared freely, but the responsibility for destruction is a single one — and it hangs off the DataPtr.

The allocator. During training the pattern of memory is extremely monotonous: allocate, release, then come back with another block of the same size. With plain malloc/free, size reuse is left entirely to the system library, which cannot serve the high-frequency pattern of “release it and immediately ask for the same size again”. So the allocator maintains its own free list: every block carves out an extra 64-byte header to record its size, and on release the block goes back into the list by size, so the next request of the same size just takes it away.

class CachingAllocator : public Allocator {
  struct Header { size_t size; };                    // 记录整块分配的尺寸
  static constexpr size_t HEADER_SIZE = 64;          // 头部 64 字节,保住对齐
  std::unordered_map<size_t, std::vector<void*>> free_blocks_;  // 尺寸桶 -> 空闲块

  DataPtr allocate(size_t nbytes) const override {
    size_t total = (nbytes + HEADER_SIZE + 63) & ~63; // 请求+头部,取整到 64 的倍数
    void* ptr = nullptr;
    auto it = free_blocks_.find(total);               // 先查桶
    if (it != free_blocks_.end() && !it->second.empty()) {
      ptr = it->second.back();                       // 有块:拿最近还回来的一块
      it->second.pop_back();
    }
    if (!ptr) ptr = alloc_aligned(total);             // 桶空:才向系统要一块
    reinterpret_cast<Header*>(ptr)->size = total;     // 实际尺寸写进头部
    return DataPtr(static_cast<char*>(ptr) + HEADER_SIZE,
                   deleter, Device(DeviceType::CPU)); // 数据指针 = 头部地址 + 64
  }

  static void deleter(void* ptr) {                    // 释放侧:不用查任何表
    char*  raw  = static_cast<char*>(ptr) - HEADER_SIZE;
    size_t size = reinterpret_cast<Header*>(raw)->size;
    instance()->free_blocks_[size].push_back(raw);    // 直接回桶,不碰系统调用
  }
};

How the buckets are divided: round up to a multiple of 64. A request for 100 bytes and a request for 124 bytes land in the same bucket, so allocations of similar size share the same set of blocks; without rounding, every size gets its own bucket, the free list becomes large and scattered, and fragmentation piles up. Buckets are only opened for sizes that have actually been used, which is why it’s an unordered_map rather than a fixed-length array.

Which way to pull a block from: back(), last in first out. A block that was just returned is most likely still hot in the CPU cache, so taking it on the next allocation is the best deal, and both push_back/pop_back are O(1).

Allocation and release are asymmetric: allocation has to look up a table (find(total)), while release looks up nothing at all — what’s written on the block header is the rounded size, subtract 64 and read it, straight back into the bucket.

The allocation policy is per device. The core code doesn’t hardcode an allocator; it looks up the allocator entry point by device type: on CPU it’s the caching allocator singleton above, under a CUDA build it’s swapped for its own implementation, and there’s also a pinned memory allocator available. The upper layers (Storage construction, clone, copy) uniformly write getAllocator(device.type()), so swapping the allocator implementation never touches the core; cross-device copies also have a unified entry point, allocating on the destination side and choosing the copy path according to both devices.

class Allocator {
 public:
  virtual DataPtr allocate(size_t nbytes) const = 0;
  virtual DataPtr allocate(size_t nbytes, const Device& device) const {
    (void)device; return allocate(nbytes);      // 单设备后端的默认实现
  }
};
Allocator* getAllocator(DeviceType t);   // 按设备类型分发
Allocator* getCPUAllocator();            // CPU 单例

The device “key”: allocation uses it to pick a policy, and later the scheduler uses it to pick a kernel. The same field, two layers of use.

Storage: data pointer, byte count, allocator, variable-length flag — wrap that in an in-house intrusive reference-counted pointer intrusive_ptr<StorageImpl> (the count is embedded in the StorageImpl), and it can “be safely shared by multiple TensorImpls”. To tell whether two Storages are the same one, you just compare the StorageImpl pointers.

intrusive_ptr is one kind of reference-counted pointer, and its count doesn’t sit beside the pointer but is embedded in the pointed-to object itself. StorageImpl inherits from a base class, and that base class stores the count. So the handle itself is a single machine word — a raw pointer; copying it once is one relaxed atomic increment on the count embedded in the object, without carrying an extra pointer to a “control block” as std::shared_ptr does, and without one more level of indirection. Strong and weak references are packed into the same 64-bit atomic word, the weak count in the high bits and the strong count in the low bits, so a single atomic operation sees both; when the count reaches zero the object deletes itself, so the object must be heap-allocated and is created with make_intrusive. TensorPlay doesn’t use an off-the-shelf library but rolls its own — because the framework reference-counts TensorImpls and StorageImpls everywhere, this thing is worth writing by hand.

Keep the distinction straight: the main ownership chain (Tensor → TensorImpl → Storage → StorageImpl) uses this in-house pointer, while auxiliary metadata like autograd_meta_ and quantizer_ still uses std::shared_ptr. The reason is simple: this metadata isn’t “shared bytes”, it’s a piece of optional extra state with a short lifetime that never participates in view aliasing; using the off-the-shelf control-block pointer is easier, and there’s no need to embed the count into those objects as well.

TensorImpl: flat and inlined

SparseTensor, QuantizedTensor and variants like these — make them subclasses, or make them fields?

The inheritance approach is object-oriented, but in practice it’s digging a hole: one set of virtual functions per variant, dynamic_cast all over the kernel code; once the header hierarchy is standing up, changing a single field drags the whole inheritance tree into a rebuild; and the overwhelming majority of tensors are neither sparse nor quantized. TensorPlay picks the flat class — all state in one struct, with sparse, quantization, and vmap all made into nullable fields, so a non-sparse tensor pays only the cost of a few null pointers for these capabilities.

class TensorImpl {                          // 一个扁平类,八组字段
  size_t          storage_offset_;          // ① 元素级偏移(视图从中间开始看)
  SizesAndStrides sizes_and_strides_;       // ② 形状与步长
  DType dtype_; Device device_;             // ③ 类型与放置
  VariableVersion version_counter_;         // ④ 原子版本计数器
  bool      inference_tensor_;              //    推理张量:整体禁用计数
  bool      is_contiguous_;
  MemoryFormat memory_format_;              // ⑤ 布局标签
  std::shared_ptr<AutogradMetaBase> autograd_meta_;  // ⑥ 梯度历史(从不拷贝)
  bool      is_view_;
  Storage   storage_;                       // ⑦ 引用计数的字节
  intrusive_ptr<TensorImpl> transform_value_;  // ⑧ vmap 批处理包装(自研指针)
  std::shared_ptr<Quantizer>  quantizer_;   //    量化器(普通共享指针)
};

Shape and strides get split out once more into a SizesAndStrides, saving the cost of new. The cost of new mostly isn’t in the system call — a small allocation goes through the allocator’s fast path, a matter of tens of nanoseconds; what really hurts comes later: every allocation has to enter the allocator’s address book to find a block, pull the block out, and write down an entry; a new block is cold, it doesn’t sit next to its neighbors, so later traversals, comparisons, and copies of sizes will most likely miss the cache; plus two allocations and two destructions. So as long as the rank is no more than 5, two fixed-length arrays are inlined into the object, and construction, copying, and destruction all bypass the heap — the extra cost is just 80 bytes on the TensorImpl body itself; past 5 it spills onto the heap, and shrinking back releases it automatically.

class SizesAndStrides {
  static constexpr size_t kInlineSize = 5;   // 秩 ≤5 全内联
  size_t  size_ = 0;                          // 秩
  int64_t inline_sizes_[5]   = {};            // 定长内联数组
  int64_t inline_strides_[5] = {};
  std::unique_ptr<int64_t[]> heap_sizes_, heap_strides_;  // 超过 5 维才落堆

  const int64_t* sizes_data() const {
    return is_heap() ? heap_sizes_.get() : inline_sizes_;  // 一次分支
  }
};

The formula for generating contiguous strides is one line: strides[i] = strides[i+1] * max(size[i+1], 1), with the last dimension’s stride being 1. The shape (2,0) comes out as (1,1), with no division by zero.

The Tensor body itself is very thin: its only data member is the in-house intrusive_ptr<TensorImpl>, plus code-generated methods.

The version counter

Why does a tensor need to know “how many times I’ve been modified”?

Because of backward: backpropagation needs the intermediate values from the forward pass. The backward of y = x.exp() needs to replay exp’s output, and that output is saved. If the user modifies it in place before backward, the framework has only two choices: use the wrong value and silently compute the wrong gradient, or refuse. TensorPlay picks the latter — record the version number when saving, and throw an exception at unpack time if they don’t line up.

Version guard timeline: how SavedVariable stops silent errors

struct VersionCounter {
  std::atomic<uint32_t> version_{0};
  void    bump()             noexcept { version_.fetch_add(1, std::memory_order_relaxed); }
  uint32_t current_version() const noexcept { return version_.load(std::memory_order_relaxed); }
};

class VariableVersion {                      // 别名之间共享同一个计数器对象
  std::shared_ptr<VersionCounter> counter_;
  bool enabled_ = true;
public:
  VariableVersion()                          // 急切分配:构造即建计数器
    : counter_(std::make_shared<VersionCounter>()) {}
  void bump() { if (!enabled_) return; counter_->bump(); }
  bool shares_counter(const VariableVersion& o) const {
    return counter_ && counter_ == o.counter_;             // 指针同一性
  }
};

It’s easy to think the make_shared in the constructor is redundant — why not create the counter on the first modification? You can’t, and the problem is views. A view must share the counter with its base tensor (modifying the view means modifying the base tensor), and that sharing works because the shared_ptr in the two VariableVersions points to the same object. With lazy allocation, whichever view is created first gets its own counter, and later an in-place modification of the base tensor leaves the view’s version absolutely motionless. The guard is still there, but it no longer guards anything — this kind of error doesn’t crash, it just occasionally hands you the wrong gradient, the hardest kind to track down. So the counter would rather pay for one wasted allocation than break the guarantee of “same source at creation”.

In inference mode the counter is disabled entirely, saving the atomic operations; but modifying an inference tensor outside inference mode throws an exception directly.

Let’s run this guard:

import tensorplay as tp

x = tp.Tensor([1.0, 2.0], requires_grad=True)
h = x.exp()         # 反向传播需要重放 h 的前向输出
print(h._version)   # 0
h.add_(1.0)         # 原地修改,版本 0 -> 1
y = h.sum()
y.backward()
# RuntimeError: one of the variables needed for gradient computation
# has been modified by an inplace operation

The copy constructor

TensorImpl’s copy constructor is the densest dozen or so lines in the whole class. It has to answer one question: for TensorImpl b = a, which things are “the data itself” and must come along, and which are “where this data came from” and must not?

The answer lives entirely in the choices made in the initializer list:

TensorImpl::TensorImpl(const TensorImpl& other)
    : storage_offset_(other.storage_offset_),
      sizes_and_strides_(other.sizes_and_strides_),   // 形状步长:拷贝
      dtype_(other.dtype_), device_(other.device_),
      version_counter_(other.version_counter_),        // 计数器:共享同一对象
      is_contiguous_(other.is_contiguous_),
      memory_format_(other.memory_format_),
      storage_(other.storage_),                        // 存储:共享(拷贝即别名)
      sparse_state_(other.sparse_state_),              // 稀疏组件:共享
      transform_value_(other.transform_value_),        // 批处理包装:共享
      quantizer_(other.quantizer_) {}                  // 量化器:共享
// autograd_meta_ 刻意不在列表里 —— 拷贝从“零历史”开始

Storage is shared, so b = a is an alias, and modifying b’s data is modifying a’s data. The version counter is shared too.

And autograd_meta_ is not in the list: the copied tensor has no gradient history, because the history describes “how this value came about”. Imagine it got carried off along the way — after b = a, b.backward() would pour the gradient into a’s parent nodes, which amounts to consuming a on the user’s behalf where he can’t see it, and without reporting any error. detach() being made into “a copy plus clearing the history” rather than a new kind of view is the same judgment: sharing a TensorImpl shares things that shouldn’t be shared.

Views

View semantics have three invariants: reuse the base tensor’s Storage, alias its version counter, and carry the quantizer along if there is one. The trouble is that view operations are a whole family — view, transpose, select, slice, permute, views across dtypes. If the invariants were written into six or seven ops, each one writing them out for itself, there’d come a day when someone adding a new view drops one of them — say, forgets to share the counter, and from then on the version guard silently fails for this kind of view, which almost never shows up in tests.

What TensorPlay does is funnel all aliasing operations into a single as_strided, so the invariants are written once:

Tensor Tensor::as_strided(const std::vector<int64_t>& size,
                          const std::vector<int64_t>& stride,
                          std::optional<int64_t> storage_offset) const {
  const int64_t offset =
      storage_offset.value_or(impl_->storage_offset());   // 缺省沿用基偏移
  if (impl_->device().is_vulkan()) {
    /* Vulkan 无步长寻址:唯一走内核物化的分支 */
  }
  Tensor out = Tensor(impl_->storage(), size, stride,
                      impl_->dtype(), offset);            // ① 复用基张量 Storage
  out.unsafeGetTensorImpl()->share_version_counter(*impl_); // ② 别名版本计数器
  if (impl_->has_quantizer())
    out.unsafeGetTensorImpl()->set_quantizer(impl_->quantizer()); // ③ 携带量化器
  return out;
}

So each concrete view is thin enough to be left with nothing but offset arithmetic, and once it has done the math it hands off to as_strided:

  • select(dim, i): new offset = storage_offset + i * stride[dim], one dimension dropped
  • slice(start, step): the offset advances start * stride[dim], the stride scales by step
  • view: infers the -1 dimension, computes a compatible stride, and errors out when the block structure doesn’t line up

Zero-copy views: two labels, one warehouse, a shared version

Vulkan is the only one that bypasses this path: data in a shader has no strided addressing, so a view can only be materialized into a real copy.

Clone is the other side of the same coin as a view, and it has only two invariants: brand-new bytes, version starting at zero. Under the default policy, a non-overlapping and dense input keeps its original strides (cloning a transpose still gives a transposed layout), otherwise it falls back to row-major; on CPU, when both sides are row-major, it goes straight to memcpy. The easiest line to miss: the reset_version at the end. The internal copy_ touches the destination’s counter, and what clone produces is a new value that has never been modified, so the version must go back to zero — otherwise the very first reconciliation would report a false alarm.

Tensor detail::clone_impl(const Tensor& self, std::optional<MemoryFormat> fmt) {
  const auto format  = fmt.value_or(MemoryFormat::Preserve);
  const auto sizes_v = static_cast<std::vector<int64_t>>(self.shape());
  std::vector<int64_t> strides;
  if (format == Preserve &&
      SizesAndStrides::is_non_overlapping_and_dense(sizes_v, self.strides()))
    strides = static_cast<std::vector<int64_t>>(self.strides()); // 转置克隆出转置
  else
    strides = SizesAndStrides::compute_contiguous_strides(sizes_v); // 落回行优先

  Storage storage(numel(self) * itemsize(self),          // 全新字节
                  getAllocator(self.device().type()), self.device());
  Tensor t(std::make_shared<TensorImpl>(storage, sizes_v, strides, self.dtype(), 0));

  if (self.device().is_cpu() && row_major(strides) && row_major(self.strides()))
    std::memcpy(t.data_ptr(), self.data_ptr(), numel(self) * itemsize(self)); // 快路径
  else
    t.copy_(self);                                       // 感知布局的拷贝

  t.unsafeGetTensorImpl()->reset_version();              // 克隆从“未被修改”开始
  return t;
}

These behaviors, measured:

import tensorplay as tp

a = tp.Tensor([[1., 2., 3.], [4., 5., 6.]])
v = a[0]                                # 视图
print(a.data_ptr() == v.data_ptr())     # True —— 同一块存储
c = a.clone()                           # 克隆
print(a.data_ptr() == c.data_ptr())     # False —— 新存储
v[0] = 99.0                             # 改视图,a 跟着变
c[0] = -1.0                             # 改克隆,a 不动
print(a)
# tensor([[99.,  2.,  3.],
#         [ 4.,  5.,  6.]])

The binding layer: how a Python object becomes a C++ Tensor

First let me clarify one thing: that line b = a at the very beginning actually does nothing in Python — b and a are the same PyObject, the reference count goes up by one, and the copy constructor at the C++ level isn’t even disturbed. The four questions above were about the copy constructor of Tensor b = a at the C++ level; what genuinely crosses the binding layer in Python are operations like a[0] and a.clone() — arguments unpacked in, results wrapped back out, all of it on this layer.

Wrapping back (C++ → Python): tpx_py_wrap takes a TensorImpl pointer as the key and looks up a cache table; on a hit it returns the Python object that already exists — the same impl will never be wrapped into two objects, so identity in Python corresponds one-to-one with identity in C++. A capsule attribute hangs off the wrapper, and when Python’s garbage collector takes it away, the capsule’s destructor wipes the cache entry on the way out, leaving no dirty reference.

Unwrapping (Python → C++): the tpx_py_tensor family comes in three forms — by value, by reference, and by mutable reference: by reference saves a pair of refcount operations, and in-place ops must go through the mutable reference, otherwise you write to a copy. Before entering the kernel there’s one overall probe; the version number and requires_grad are read straight off the C++ side in a single call, and if it isn’t a legitimate wrapper it returns -1 without throwing.

Entry points: dual paths. The pybind11 bridge plus a code-generated FASTCALL entry — the argument array goes straight through with no temporary tuple; when registering the generated table it checks for duplicates first and simply skips names pybind11 has already installed. Argument parsing is done inside the generated entry: for multiple overloads, a non-throwing type probe picks the candidate first (16 types plus an optional flag), and only after a hit does the real check happen. The GIL is released while the kernel executes, so large matrix operations don’t sit holding the interpreter.

Where these trade-offs land

Every multiple-choice question above comes down to the same place: the reference count is embedded in the pointed-to object; the price is writing your own pointer type, and what you buy is a handle that occupies only a single machine word.

  • Flat implementation class: TensorImpl doesn’t take the inheritance tree. Sparse and quantized aren’t subclasses but a few null pointers on the same class — layout-independent sparse tensors and quantized tensors pay their cost only when they’re actually used, and on the dense path the only extra is one null check on the pointer.
  • Counting merged into one word: intrusive_ptr puts the strong count in the low bits and the weak count in the high bits and merges them into a single 64-bit atomic word, so one atomic operation sees both at once. Copying a handle is one relaxed atomic increment, the handle itself is one machine word, there’s no second machine word that “holds a separate copy of the count”, and there’s no extra indirection through a control block. The price is that make_intrusive requires heap allocation — that’s the trade written down in p10/include/IntrusivePtr.h.
  • The binding layer: the Python entry is layered on top of pybind11, but the hot path goes through the generated FASTCALL entry — the argument array goes straight through, wrapping back has an identity cache, unwrapping has a non-throwing type probe, and it tries not to create intermediate objects at the Python layer.

References