走进 TensorPlay（一）：一个 Tensor 的背后

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; }
};

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);    // 直接回桶，不碰系统调用
  }
};

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 单例

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_;   //    量化器（普通共享指针）
};

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_;  // 一次分支
  }
};

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_;             // 指针同一性
  }
};

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

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_ 刻意不在列表里 —— 拷贝从“零历史”开始

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;
}

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;
}

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.]])