PyTorch Tensor核心机制深度剖析：从存储到视图的完整实现

 1
 2
 3
 4
 5
 6
 7
 8
 9
10
11
12
13

┌─────────────────────────────────────────┐
│           Python torch.Tensor           │ ← Python用户接口
├─────────────────────────────────────────┤
│          TensorBase C++ 接口            │ ← C++基础接口
├─────────────────────────────────────────┤
│           TensorImpl 核心               │ ← 张量实现核心
├─────────────────────────────────────────┤
│            Storage 存储层               │ ← 数据存储管理
├─────────────────────────────────────────┤
│          StorageImpl 存储实现           │ ← 底层存储实现
├─────────────────────────────────────────┤
│            DataPtr 数据指针             │ ← 内存数据指针
└─────────────────────────────────────────┘

 1
 2
 3
 4
 5
 6
 7
 8
 9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66

class C10_API TensorImpl : public c10::intrusive_ptr_target {
 private:
  // 1. 存储相关
  Storage storage_;                    // 底层数据存储
  
  // 2. 形状和步长信息  
  c10::impl::SizesAndStrides sizes_and_strides_;  // 尺寸和步长
  SymInt storage_offset_;              // 存储偏移量
  
  // 3. 类型和设备信息
  ScalarType dtype_;                   // 数据类型（float32、int64等）
  Device device_;                      // 设备信息（CPU、CUDA等）
  
  // 4. 分发和优化相关
  DispatchKeySet key_set_;             // 分发键集合
  MemoryFormat memory_format_;         // 内存格式（连续、NHWC等）
  
  // 5. 梯度和版本管理
  std::unique_ptr<c10::impl::PyObjectSlot> pyobj_slot_;  // Python对象槽
  c10::impl::NamedTensorMetaInterface* named_tensor_meta_; // 命名张量元数据
  c10::VariableVersion version_counter_;  // 版本计数器
  c10::impl::AutogradMetaInterface* autograd_meta_;  // 自动微分元数据

 public:
  // 构造函数 - 创建具有给定存储和选项的张量
  TensorImpl(
      Storage&& storage,
      DispatchKeySet key_set,
      const caffe2::TypeMeta& data_type
  );
  
  // 获取底层存储
  const Storage& storage() const { return storage_; }
  
  // 获取张量维度
  int64_t dim() const { return sizes_and_strides_.size(); }
  
  // 获取形状信息
  IntArrayRef sizes() const {
    return IntArrayRef{
        sizes_and_strides_.sizes_data(), 
        sizes_and_strides_.size()
    };
  }
  
  // 获取步长信息
  IntArrayRef strides() const {
    return IntArrayRef{
        sizes_and_strides_.strides_data(),
        sizes_and_strides_.size()
    };
  }
  
  // 获取元素总数
  int64_t numel() const {
    return c10::size_to_dim_(dim(), sizes());
  }
  
  // 检查张量是否连续
  bool is_contiguous(at::MemoryFormat memory_format = at::MemoryFormat::Contiguous) const;
  
  // 调整张量大小
  virtual void resize_(IntArrayRef size, std::optional<IntArrayRef> stride = std::nullopt);
  
  // 其他核心方法...
};

 1
 2
 3
 4
 5
 6
 7
 8
 9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50

namespace c10::impl {

class SizesAndStrides {
 private:
  // 小张量优化：直接在对象内存储而非堆分配
  static constexpr size_t kInlineStorageSize = 5;
  
  // 联合体优化内存使用
  union {
    int64_t inline_storage_[kInlineStorageSize];  // 内联存储
    int64_t* heap_storage_;                       // 堆存储指针
  };
  
  uint8_t size_;           // 维度数量
  bool is_heap_allocated_; // 是否使用堆分配

 public:
  // 构造函数
  explicit SizesAndStrides(size_t size) : size_(size) {
    if (C10_LIKELY(size <= kInlineStorageSize)) {
      // 小张量使用栈内存，性能更好
      is_heap_allocated_ = false;
    } else {
      // 大张量使用堆内存
      is_heap_allocated_ = true;
      heap_storage_ = static_cast<int64_t*>(
          std::malloc(sizeof(int64_t) * size * 2)  // size + strides
      );
    }
  }
  
  // 获取形状数据指针
  const int64_t* sizes_data() const noexcept {
    return is_heap_allocated_ ? heap_storage_ : inline_storage_;
  }
  
  // 获取步长数据指针  
  const int64_t* strides_data() const noexcept {
    const int64_t* base = sizes_data();
    return base + size_;  // 步长紧跟在形状之后
  }
  
  // 设置形状信息
  void set_sizes(IntArrayRef newSizes) {
    TORCH_CHECK(newSizes.size() == size_);
    std::copy(newSizes.begin(), newSizes.end(), sizes_data_non_const());
  }
};

} // namespace c10::impl

 1
 2
 3
 4
 5
 6
 7
 8
 9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52

namespace c10 {

// Storage 接口类 - 提供统一的存储访问接口
struct C10_API Storage {
 private:
  c10::intrusive_ptr<StorageImpl> storage_impl_;  // 指向具体实现
  
 public:
  // 默认构造函数 - 创建空存储
  Storage() = default;
  
  // 从StorageImpl构造
  Storage(c10::intrusive_ptr<StorageImpl> ptr) 
      : storage_impl_(std::move(ptr)) {}
  
  // 使用字节大小创建存储
  Storage(
      use_byte_size_t /*use_byte_size*/,
      const SymInt& size_bytes,
      Allocator* allocator = nullptr,
      bool resizable = false
  ) : storage_impl_(c10::make_intrusive<StorageImpl>(
          StorageImpl::use_byte_size_t(),
          size_bytes,
          allocator,
          resizable)) {}
  
  // 使用预分配内存创建存储  
  Storage(
      use_byte_size_t /*use_byte_size*/,
      size_t size_bytes,
      at::DataPtr data_ptr,
      at::Allocator* allocator = nullptr,
      bool resizable = false
  );
  
  // 核心访问方法
  void* mutable_data() const { return storage_impl_->mutable_data(); }
  const void* data() const { return storage_impl_->data(); }
  
  SymInt nbytes() const { return storage_impl_->nbytes(); }
  Device device() const { return storage_impl_->device(); }
  
  bool resizable() const { return storage_impl_->resizable(); }
  Allocator* allocator() { return storage_impl_->allocator(); }
  
  // 引用计数管理
  bool unique() const { return storage_impl_.use_count() == 1; }
  int use_count() const { return storage_impl_.use_count(); }
};

} // namespace c10

  1
  2
  3
  4
  5
  6
  7
  8
  9
 10
 11
 12
 13
 14
 15
 16
 17
 18
 19
 20
 21
 22
 23
 24
 25
 26
 27
 28
 29
 30
 31
 32
 33
 34
 35
 36
 37
 38
 39
 40
 41
 42
 43
 44
 45
 46
 47
 48
 49
 50
 51
 52
 53
 54
 55
 56
 57
 58
 59
 60
 61
 62
 63
 64
 65
 66
 67
 68
 69
 70
 71
 72
 73
 74
 75
 76
 77
 78
 79
 80
 81
 82
 83
 84
 85
 86
 87
 88
 89
 90
 91
 92
 93
 94
 95
 96
 97
 98
 99
100
101
102
103
104

struct C10_API StorageImpl : public c10::intrusive_ptr_target {
 private:
  // 核心数据成员
  DataPtr data_ptr_;                    // 数据指针（包含删除器）
  SymInt size_bytes_;                   // 字节大小（支持符号整数）
  bool size_bytes_is_heap_allocated_;   // 大小是否堆分配
  bool resizable_;                      // 是否可调整大小
  bool received_cuda_;                  // 是否接收过CUDA数据
  Allocator* allocator_;                // 内存分配器
  
  // 扩展元数据（用于错误处理等）
  std::unique_ptr<StorageExtraMeta> extra_meta_;
  
 public:
  // 构造函数 - 使用字节大小和数据指针
  StorageImpl(
      use_byte_size_t /*use_byte_size*/,
      SymInt size_bytes,
      at::DataPtr data_ptr,
      at::Allocator* allocator,
      bool resizable
  ) : data_ptr_(std::move(data_ptr)),
      size_bytes_(std::move(size_bytes)),
      size_bytes_is_heap_allocated_(size_bytes_.is_heap_allocated()),
      resizable_(resizable),
      received_cuda_(false),
      allocator_(allocator) {
    
    // 可调整大小的存储必须有分配器
    if (resizable) {
      TORCH_INTERNAL_ASSERT(
          allocator_, 
          "For resizable storage, allocator must be provided"
      );
    }
    refresh_has_data_ptr_check();
  }
  
  // 获取原始数据指针
  void* data() {
    if (!has_data_ptr_check_) [[unlikely]] {
      check_data_ptr();
    }
    return data_ptr_.get();
  }
  
  // 获取字节大小
  SymInt nbytes() const {
    return size_bytes_;
  }
  
  // 调整存储大小
  void resize_(const SymInt& new_size_bytes) {
    TORCH_CHECK(resizable(), "Trying to resize storage that is not resizable");
    
    if (new_size_bytes == 0) {
      // 大小为0时释放内存但保留分配器
      data_ptr_ = allocator_->allocate(0);
    } else {
      // 重新分配内存
      DataPtr new_data = allocator_->allocate(new_size_bytes.as_int_unchecked());
      
      // 如果有旧数据，拷贝到新内存（取较小值）
      if (data_ptr_) {
        auto old_size = size_bytes_.as_int_unchecked();
        auto copy_size = std::min(old_size, new_size_bytes.as_int_unchecked());
        if (copy_size > 0) {
          std::memcpy(new_data.get(), data_ptr_.get(), copy_size);
        }
      }
      
      data_ptr_ = std::move(new_data);
    }
    
    size_bytes_ = new_size_bytes;
    refresh_has_data_ptr_check();
  }
  
  // 设备信息
  Device device() const {
    return data_ptr_.device();
  }
  
  // 分配器信息
  Allocator* allocator() {
    return allocator_;
  }
  
 private:
  // 检查数据指针有效性
  void check_data_ptr() {
    if (extra_meta_ && extra_meta_->custom_data_ptr_error_msg_) {
      TORCH_CHECK(false, extra_meta_->custom_data_ptr_error_msg_.value());
    }
    if (!data_ptr_) {
      throwNullDataPtrError();
    }
  }
  
  // 刷新数据指针检查状态
  void refresh_has_data_ptr_check() {
    has_data_ptr_check_ = (!extra_meta_ || !extra_meta_->custom_data_ptr_error_msg_) && data_ptr_;
  }
};

 1
 2
 3
 4
 5
 6
 7
 8
 9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56

namespace at {

// DataPtr - 设备感知的智能指针
class DataPtr {
 private:
  void* ptr_;                    // 原始指针
  std::unique_ptr<void, DeleterFnPtr> ctx_;  // 删除器上下文
  Device device_;                // 设备信息
  
 public:
  // 构造函数
  DataPtr(void* data, Device device) 
      : ptr_(data), device_(device) {}
      
  DataPtr(void* data, void* ctx, DeleterFnPtr deleter, Device device)
      : ptr_(data), 
        ctx_(ctx, deleter),
        device_(device) {}
  
  // 移动构造和赋值
  DataPtr(DataPtr&& other) noexcept
      : ptr_(other.ptr_), 
        ctx_(std::move(other.ctx_)), 
        device_(other.device_) {
    other.ptr_ = nullptr;
  }
  
  // 获取原始指针
  void* get() const { return ptr_; }
  
  // 释放所有权
  void* release() noexcept {
    ctx_.release();
    auto result = ptr_;
    ptr_ = nullptr;
    return result;
  }
  
  // 获取设备信息
  Device device() const { return device_; }
  
  // 重置指针
  void clear() {
    ctx_.reset();
    ptr_ = nullptr;
  }
  
  // 类型转换
  template<typename T>
  T* get() const { return static_cast<T*>(ptr_); }
  
  // 布尔转换
  explicit operator bool() const { return ptr_ != nullptr; }
};

} // namespace at

 1
 2
 3
 4
 5
 6
 7
 8
 9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96

namespace c10 {

// 抽象分配器接口
class C10_API Allocator {
 public:
  virtual ~Allocator() = default;
  
  // 分配内存
  virtual DataPtr allocate(size_t n) const = 0;
  
  // 获取删除器函数指针
  virtual DeleterFnPtr raw_deleter() const { 
    return nullptr; 
  }
  
  // 复制数据（可选优化）
  virtual void copy_data(void* dest, const void* src, std::size_t count) const {
    std::memcpy(dest, src, count);
  }
};

// CPU默认分配器
class DefaultCPUAllocator : public Allocator {
 public:
  DataPtr allocate(size_t nbytes) const override {
    void* data = aligned_alloc(gAlignment, nbytes);
    if (!data) {
      TORCH_CHECK(false, "DefaultCPUAllocator: can't allocate memory");
    }
    return {data, data, &free, kCPU};
  }
  
  DeleterFnPtr raw_deleter() const override {
    return &free;
  }
};

// CPU缓存分配器 - 减少分配/释放开销
class CPUCachingAllocator : public Allocator {
 private:
  // 内存块结构
  struct Block {
    void* ptr;
    size_t size;
    bool in_use;
  };
  
  mutable std::mutex mutex_;
  mutable std::vector<Block> cached_blocks_;
  
 public:
  DataPtr allocate(size_t nbytes) const override {
    std::lock_guard<std::mutex> lock(mutex_);
    
    // 1. 尝试从缓存中找到合适的块
    auto it = std::find_if(cached_blocks_.begin(), cached_blocks_.end(),
        [nbytes](const Block& block) {
          return !block.in_use && block.size >= nbytes;
        });
    
    if (it != cached_blocks_.end()) {
      // 找到缓存块，标记为使用中
      it->in_use = true;
      return {it->ptr, this, &CPUCachingAllocator::free_cached, kCPU};
    }
    
    // 2. 没有合适的缓存，分配新内存
    void* ptr = aligned_alloc(gAlignment, nbytes);
    if (!ptr) {
      TORCH_CHECK(false, "CPUCachingAllocator: allocation failed");
    }
    
    // 添加到缓存中
    cached_blocks_.push_back({ptr, nbytes, true});
    
    return {ptr, this, &CPUCachingAllocator::free_cached, kCPU};
  }
  
  // 释放内存到缓存
  static void free_cached(void* ctx, void* ptr) {
    auto* allocator = static_cast<CPUCachingAllocator*>(ctx);
    std::lock_guard<std::mutex> lock(allocator->mutex_);
    
    auto it = std::find_if(allocator->cached_blocks_.begin(), 
                          allocator->cached_blocks_.end(),
        [ptr](const Block& block) {
          return block.ptr == ptr;
        });
    
    if (it != allocator->cached_blocks_.end()) {
      it->in_use = false;  // 标记为可用，不实际释放
    }
  }
};

} // namespace c10

 1
 2
 3
 4
 5
 6
 7
 8
 9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75

// 视图创建的核心逻辑
Tensor TensorImpl::create_view(
    IntArrayRef size,
    IntArrayRef stride, 
    int64_t storage_offset
) {
  // 1. 创建新的TensorImpl，但共享Storage
  auto impl = c10::make_intrusive<TensorImpl>(
      Storage(storage()),           // 共享底层存储
      key_set(),                   // 继承分发键
      dtype()                      // 继承数据类型
  );
  
  // 2. 设置新的形状和步长
  impl->set_sizes_and_strides(size, stride);
  impl->set_storage_offset(storage_offset);
  
  // 3. 继承其他属性
  impl->set_device(device());
  impl->set_layout(layout());
  
  // 4. 返回包装的Tensor对象
  return at::detail::make_tensor<TensorImpl>(std::move(impl));
}

// 常见视图操作的实现
namespace at::native {

// 切片操作 - tensor[start:end]
Tensor slice(const Tensor& self, int64_t dim, int64_t start, int64_t end, int64_t step) {
  // 计算新的形状
  auto sizes = self.sizes().vec();
  sizes[dim] = (end - start + step - 1) / step;
  
  // 计算新的步长
  auto strides = self.strides().vec(); 
  strides[dim] *= step;
  
  // 计算存储偏移
  auto offset = self.storage_offset() + start * self.stride(dim);
  
  // 创建视图
  return self.as_strided(sizes, strides, offset);
}

// 转置操作 - tensor.t()  
Tensor transpose(const Tensor& self, int64_t dim0, int64_t dim1) {
  auto sizes = self.sizes().vec();
  auto strides = self.strides().vec();
  
  // 交换指定维度的尺寸和步长
  std::swap(sizes[dim0], sizes[dim1]);
  std::swap(strides[dim0], strides[dim1]);
  
  // 创建视图（共享存储，不同形状）
  return self.as_strided(sizes, strides, self.storage_offset());
}

// reshape操作 - 改变形状
Tensor reshape(const Tensor& self, IntArrayRef shape) {
  // 1. 检查元素总数是否匹配
  auto numel = c10::multiply_integers(shape);
  TORCH_CHECK(numel == self.numel(), "reshape: incompatible shape");
  
  // 2. 如果可能，尝试创建视图
  auto stride = compute_stride_for_reshape(self, shape);
  if (stride.has_value()) {
    return self.as_strided(shape, *stride, self.storage_offset());
  }
  
  // 3. 否则，需要拷贝数据并重新排列
  return self.contiguous().view(shape);
}

} // namespace at::native

 1
 2
 3
 4
 5
 6
 7
 8
 9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56

namespace c10 {

// 计算连续张量的默认步长
std::vector<int64_t> compute_contiguous_strides(IntArrayRef sizes) {
  std::vector<int64_t> strides(sizes.size());
  
  // 从最后一维开始，步长为1
  strides.back() = 1;
  
  // 向前计算每一维的步长
  for (int i = static_cast<int>(sizes.size()) - 2; i >= 0; --i) {
    strides[i] = strides[i + 1] * sizes[i + 1];
  }
  
  return strides;
}

// 示例：形状为[2,3,4]的张量
// sizes = [2, 3, 4]
// strides = [12, 4, 1]  // 12=3*4, 4=4*1, 1
// 元素[i,j,k]在内存中的位置 = i*12 + j*4 + k*1

// 计算NHWC格式的步长（通道优先）
std::vector<int64_t> compute_channels_last_strides(IntArrayRef sizes) {
  TORCH_CHECK(sizes.size() == 4, "channels_last only supports 4D tensors");
  
  auto N = sizes[0], C = sizes[1], H = sizes[2], W = sizes[3];
  
  // NHWC布局：步长为[H*W*C, 1, W*C, C]
  return {H * W * C, 1, W * C, C};
}

// 检查张量是否连续
bool is_contiguous_with_format(
    IntArrayRef sizes,
    IntArrayRef strides,
    MemoryFormat memory_format
) {
  switch (memory_format) {
    case MemoryFormat::Contiguous: {
      auto expected_strides = compute_contiguous_strides(sizes);
      return strides == expected_strides;
    }
    
    case MemoryFormat::ChannelsLast: {
      if (sizes.size() != 4) return false;
      auto expected_strides = compute_channels_last_strides(sizes);
      return strides == expected_strides;
    }
    
    default:
      return false;
  }
}

} // namespace c10

 1
 2
 3
 4
 5
 6
 7
 8
 9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38

namespace at {

// 连续性检查的核心实现（来源：PyTorch源码剖析）
bool TensorBase::is_contiguous_or_false(MemoryFormat memory_format) const {
  if (impl_->is_contiguous_default()) {
    // 快速路径：如果已知是默认连续的
    return memory_format == MemoryFormat::Contiguous;
  }
  
  // 详细检查连续性
  return impl_->is_contiguous_memory_format(memory_format);
}

// 高效的连续化操作
TensorBase TensorBase::contiguous(MemoryFormat memory_format) const {
  if (is_contiguous_or_false(memory_format)) {
    // 已经连续，直接返回（避免不必要的拷贝）
    return *this;
  } else {
    // 需要重新排列内存布局
    return __dispatch_contiguous(memory_format);
  }
}

// expect_contiguous：性能优化的连续性检查
c10::MaybeOwned<TensorBase> TensorBase::expect_contiguous(
    MemoryFormat memory_format) const & {
  
  if (is_contiguous_or_false(memory_format)) {
    // 已经连续，返回借用的引用（避免引用计数开销）
    return c10::MaybeOwned<TensorBase>::borrowed(*this);
  } else {
    // 需要创建连续版本，返回拥有的副本
    return c10::MaybeOwned<TensorBase>::owned(__dispatch_contiguous(memory_format));
  }
}

} // namespace at

 1
 2
 3
 4
 5
 6
 7
 8
 9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85

namespace c10 {

// stride计算的完整数学原理（基于源码分析总结）
class StrideCalculator {
 public:
  // 计算连续张量的步长
  static std::vector<int64_t> contiguous_strides(IntArrayRef sizes) {
    /*
     * 连续张量的步长计算公式：
     * stride[i] = ∏(j=i+1 to n-1) size[j]
     * 
     * 例如：shape=[2,3,4]
     * stride[2] = 1                    // 最后一维步长为1
     * stride[1] = size[2] = 4          // 倒数第二维
     * stride[0] = size[1]*size[2] = 12 // 第一维
     */
    
    const int64_t ndim = sizes.size();
    if (ndim == 0) return {};
    
    std::vector<int64_t> strides(ndim);
    strides[ndim - 1] = 1;  // 最后一维步长为1
    
    // 从后向前计算步长
    for (int64_t i = ndim - 2; i >= 0; --i) {
      strides[i] = strides[i + 1] * sizes[i + 1];
    }
    
    return strides;
  }
  
  // 计算通道最后格式(NHWC)的步长
  static std::vector<int64_t> channels_last_strides(IntArrayRef sizes) {
    /*
     * NHWC格式的步长计算：
     * - N(batch): stride = H * W * C
     * - H(height): stride = W * C  
     * - W(width): stride = C
     * - C(channel): stride = 1
     */
    
    TORCH_CHECK(sizes.size() == 4, "channels_last requires 4D tensor");
    
    const int64_t N = sizes[0], C = sizes[1], H = sizes[2], W = sizes[3];
    return {H * W * C, 1, W * C, C};
  }
  
  // 计算3D通道最后格式的步长
  static std::vector<int64_t> channels_last_3d_strides(IntArrayRef sizes) {
    /*
     * NDHWC格式的步长计算：
     * - N: stride = D * H * W * C
     * - D: stride = H * W * C
     * - H: stride = W * C
     * - W: stride = C
     * - C: stride = 1
     */
    
    TORCH_CHECK(sizes.size() == 5, "channels_last_3d requires 5D tensor");
    
    const int64_t N = sizes[0], C = sizes[1], D = sizes[2], H = sizes[3], W = sizes[4];
    return {D * H * W * C, 1, H * W * C, W * C, C};
  }
  
  // 检查步长是否有效（基于网上分析的边界情况处理）
  static bool are_strides_valid(IntArrayRef sizes, IntArrayRef strides) {
    if (sizes.size() != strides.size()) return false;
    
    // 检查步长的数学约束
    for (size_t i = 0; i < sizes.size(); ++i) {
      // 大小为0或1的维度，步长可以是任意值
      if (sizes[i] <= 1) continue;
      
      // 步长不能为负数
      if (strides[i] < 0) return false;
      
      // 步长为0时，该维度大小必须为1
      if (strides[i] == 0 && sizes[i] > 1) return false;
    }
    
    return true;
  }
};

} // namespace c10

 1
 2
 3
 4
 5
 6
 7
 8
 9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69

namespace at {

// TensorAccessor - 提供类型安全的多维数组访问
template<typename T, size_t N, template<typename U> class PtrTraits = DefaultPtrTraits>
class TensorAccessor {
 private:
  T* data_;                     // 数据指针
  int64_t sizes_[N];           // 各维大小
  int64_t strides_[N];         // 各维步长
  
 public:
  TensorAccessor(T* data, IntArrayRef sizes, IntArrayRef strides) 
      : data_(data) {
    TORCH_CHECK(sizes.size() == N && strides.size() == N);
    std::copy(sizes.begin(), sizes.end(), sizes_);
    std::copy(strides.begin(), strides.end(), strides_);
  }
  
  // 多维访问操作符
  template<typename... Args>
  T& operator()(Args... indices) {
    static_assert(sizeof...(Args) == N, "Wrong number of indices");
    return data_[compute_offset(indices...)];
  }
  
  // 单维访问（递归返回低维Accessor）
  TensorAccessor<T, N-1, PtrTraits> operator[](int64_t i) {
    return TensorAccessor<T, N-1, PtrTraits>(
        data_ + i * strides_[0],
        sizes_ + 1,
        strides_ + 1
    );
  }
  
 private:
  // 计算线性偏移
  template<typename... Args>
  int64_t compute_offset(int64_t first, Args... rest) {
    return first * strides_[N - sizeof...(Args) - 1] + 
           compute_offset_impl<sizeof...(Args)>(rest...);
  }
  
  template<size_t remaining, typename... Args>
  int64_t compute_offset_impl(int64_t first, Args... rest) {
    if constexpr (remaining == 1) {
      return first * strides_[N - 1];
    } else {
      return first * strides_[N - remaining] + 
             compute_offset_impl<remaining - 1>(rest...);
    }
  }
};

// 使用示例
void example_tensor_access() {
  auto tensor = torch::rand({3, 4, 5});
  
  // 获取类型安全的访问器
  auto accessor = tensor.accessor<float, 3>();
  
  // 直接多维访问
  float value = accessor(1, 2, 3);  // 等价于 tensor[1][2][3]
  
  // 层次化访问
  auto slice = accessor[1];         // 得到2D accessor
  float value2 = slice(2, 3);       // 等价于 tensor[1][2][3]
}

} // namespace at

 1
 2
 3
 4
 5
 6
 7
 8
 9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64

// 张量创建的完整流程
namespace torch {

// Python API: torch.tensor([1, 2, 3])
Tensor tensor_from_data(py::object data, py::object dtype_obj, py::object device_obj) {
  // 1. 解析Python参数
  auto dtype = dtype_obj.is_none() ? infer_dtype(data) : to_scalar_type(dtype_obj);
  auto device = device_obj.is_none() ? Device(kCPU) : to_device(device_obj);
  
  // 2. 解析数据结构，推断形状
  auto parsed = parse_python_scalars_and_tensors(data);
  auto sizes = parsed.sizes;
  auto scalar_list = parsed.scalars;
  
  // 3. 计算总元素数
  auto numel = c10::multiply_integers(sizes);
  
  // 4. 分配存储空间
  auto dtype_size = scalarTypeToTypeMeta(dtype).itemsize();
  auto nbytes = numel * dtype_size;
  
  auto allocator = GetAllocator(device.type());
  auto storage = c10::make_intrusive<c10::StorageImpl>(
      c10::StorageImpl::use_byte_size_t(),
      nbytes,
      allocator->allocate(nbytes),
      allocator,
      /*resizable=*/true
  );
  
  // 5. 创建TensorImpl
  auto tensor_impl = c10::make_intrusive<c10::TensorImpl>(
      c10::Storage(storage),
      compute_dispatch_key_set(dtype, device),
      scalarTypeToTypeMeta(dtype)
  );
  
  // 6. 设置形状信息
  tensor_impl->set_sizes_and_strides(sizes, compute_contiguous_strides(sizes));
  
  // 7. 填充数据
  fill_tensor_from_scalars(tensor_impl.get(), scalar_list, dtype);
  
  // 8. 返回包装的Tensor对象
  return at::detail::make_tensor<c10::TensorImpl>(std::move(tensor_impl));
}

// 数据填充函数
void fill_tensor_from_scalars(c10::TensorImpl* impl, 
                             const std::vector<at::Scalar>& scalars, 
                             c10::ScalarType dtype) {
  auto data_ptr = impl->storage().mutable_data();
  
  // 根据数据类型进行特化处理
  AT_DISPATCH_ALL_TYPES_AND_COMPLEX_AND3(
      c10::kHalf, c10::kBFloat16, c10::kBool, dtype, "fill_from_scalars", [&] {
        auto typed_data = static_cast<scalar_t*>(data_ptr);
        for (size_t i = 0; i < scalars.size(); ++i) {
          typed_data[i] = scalars[i].to<scalar_t>();
        }
      });
}

} // namespace torch

 1
 2
 3
 4
 5
 6
 7
 8
 9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44

namespace c10::impl {

// COW（Copy-on-Write）上下文
class COWDeleterContext {
 private:
  std::shared_ptr<void> original_data_;  // 原始数据的共享指针
  size_t original_size_;                 // 原始数据大小
  Device device_;                        // 设备信息
  
 public:
  COWDeleterContext(std::shared_ptr<void> data, size_t size, Device device)
      : original_data_(std::move(data)), original_size_(size), device_(device) {}
  
  // 当需要修改数据时，执行实际拷贝
  static void deleter_function(void* ctx, void* data_ptr) {
    auto* cow_ctx = static_cast<COWDeleterContext*>(ctx);
    
    // 如果是唯一引用，直接释放
    if (cow_ctx->original_data_.use_count() == 1) {
      // 正常释放逻辑
      GetAllocator(cow_ctx->device_.type())->raw_deleter()(data_ptr);
    }
    
    delete cow_ctx;
  }
  
  // 检查是否需要拷贝
  bool needs_cow() const {
    return original_data_.use_count() > 1;
  }
  
  // 执行写时拷贝
  DataPtr make_cow_copy() {
    auto allocator = GetAllocator(device_.type());
    auto new_data = allocator->allocate(original_size_);
    
    // 拷贝原始数据
    std::memcpy(new_data.get(), original_data_.get(), original_size_);
    
    return new_data;
  }
};

} // namespace c10::impl

  1
  2
  3
  4
  5
  6
  7
  8
  9
 10
 11
 12
 13
 14
 15
 16
 17
 18
 19
 20
 21
 22
 23
 24
 25
 26
 27
 28
 29
 30
 31
 32
 33
 34
 35
 36
 37
 38
 39
 40
 41
 42
 43
 44
 45
 46
 47
 48
 49
 50
 51
 52
 53
 54
 55
 56
 57
 58
 59
 60
 61
 62
 63
 64
 65
 66
 67
 68
 69
 70
 71
 72
 73
 74
 75
 76
 77
 78
 79
 80
 81
 82
 83
 84
 85
 86
 87
 88
 89
 90
 91
 92
 93
 94
 95
 96
 97
 98
 99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
187
188
189
190
191
192
193
194
195
196
197
198
199
200
201
202
203
204
205
206
207
208
209
210
211
212
213
214
215
216
217
218
219
220
221
222
223
224
225
226
227
228
229
230
231
232
233
234
235
236
237
238
239
240
241
242
243
244
245
246
247
248
249
250
251
252
253
254
255
256
257
258
259
260
261
262
263
264
265
266
267
268
269
270
271
272
273
274
275
276
277
278
279
280
281
282
283
284
285
286
287
288
289
290
291
292
293
294
295
296
297
298
299
300
301
302
303
304
305
306
307
308

namespace c10::cuda::CUDACachingAllocator {

// CUDA缓存分配器的核心数据结构（源码深度分析）
struct Block {
  size_t size;                    // 块大小
  size_t requested_size;          // 实际请求大小
  void* ptr;                      // 数据指针
  cudaStream_t stream;            // 关联的CUDA流
  bool allocated;                 // 是否已分配
  bool active;                    // 是否活跃使用
  c10::DeviceIndex device;        // 设备索引
  
  // 内存池相关
  MempoolId_t owner_private_pool_id;  // 私有池ID
  
  // 分割和合并支持
  Block* prev;                    // 前一个块（用于合并）
  Block* next;                    // 后一个块（用于合并）
  
  // 流相关的引用计数
  std::unordered_set<c10::cuda::CUDAStream> stream_uses;
  
  // 垃圾回收计数器
  int32_t gc_counter;
  
  // 构造函数
  Block(
      size_t size,
      size_t requested_size, 
      void* ptr,
      cudaStream_t stream,
      c10::DeviceIndex device)
      : size(size),
        requested_size(requested_size),
        ptr(ptr),
        stream(stream),
        allocated(true),
        active(true),
        device(device),
        prev(nullptr),
        next(nullptr),
        gc_counter(0) {}
};

// 内存段管理（源码分析：大块内存的管理策略）
struct Segment {
  c10::DeviceIndex device;        // 设备索引
  size_t address;                 // 内存地址
  size_t total_size;              // 总大小
  size_t allocated_size;          // 已分配大小
  size_t active_size;             // 活跃使用大小
  cudaStream_t stream;            // 创建时的流
  bool is_large;                  // 是否为大段
  bool is_expandable;             // 是否可扩展
  
  std::vector<Block*> blocks;     // 包含的块列表
  
  // 查找可用的块
  Block* find_available_block(size_t requested_size) {
    for (auto* block : blocks) {
      if (!block->allocated && block->size >= requested_size) {
        return block;
      }
    }
    return nullptr;
  }
  
  // 分割块
  Block* split_block(Block* block, size_t size) {
    if (block->size <= size) {
      return block;  // 无需分割
    }
    
    // 创建新块用于剩余空间
    size_t remaining_size = block->size - size;
    void* remaining_ptr = static_cast<char*>(block->ptr) + size;
    
    auto* new_block = new Block(
        remaining_size,
        0,  // requested_size为0表示这是分割产生的块
        remaining_ptr,
        block->stream,
        block->device
    );
    
    new_block->allocated = false;
    new_block->active = false;
    
    // 更新原块大小
    block->size = size;
    
    // 插入到块列表中
    blocks.insert(
        std::find(blocks.begin(), blocks.end(), block) + 1,
        new_block
    );
    
    return block;
  }
  
  // 合并相邻的空闲块
  void merge_free_blocks() {
    for (size_t i = 0; i < blocks.size() - 1; ++i) {
      auto* current = blocks[i];
      auto* next = blocks[i + 1];
      
      if (!current->allocated && !next->allocated) {
        // 检查是否相邻
        if (static_cast<char*>(current->ptr) + current->size == next->ptr) {
          // 合并块
          current->size += next->size;
          
          // 移除下一个块
          blocks.erase(blocks.begin() + i + 1);
          delete next;
          
          // 重新检查当前位置
          --i;
        }
      }
    }
  }
};

// 缓存分配器的核心算法（深度源码分析）
class CachingAllocatorImpl {
 private:
  // 设备特定的分配器状态
  struct DeviceAllocatorState {
    // 分大小范围的空闲块列表
    std::map<size_t, std::set<Block*>> large_free_blocks;    // 大块 (>1MB)
    std::map<size_t, std::set<Block*>> small_free_blocks;    // 小块 (<=1MB)
    
    // 活跃块列表
    std::unordered_set<Block*> active_blocks;
    
    // 内存段列表
    std::vector<std::unique_ptr<Segment>> segments;
    
    // 统计信息
    size_t allocated_size = 0;
    size_t reserved_size = 0;
    size_t peak_allocated_size = 0;
    size_t peak_reserved_size = 0;
    
    // 垃圾回收相关
    std::atomic<size_t> garbage_collect_threshold{1024 * 1024 * 1024};  // 1GB
  };
  
  // 每个设备的状态
  std::vector<DeviceAllocatorState> device_allocator_states_;
  
  // 全局锁
  mutable std::recursive_mutex mutex_;
  
 public:
  // 分配内存的核心算法
  DataPtr allocate(size_t requested_size) {
    // 对齐大小到最近的512字节边界（GPU内存对齐优化）
    size_t size = round_size(requested_size);
    
    auto device = c10::cuda::current_device();
    std::lock_guard<std::recursive_mutex> lock(mutex_);
    
    auto& state = device_allocator_states_[device];
    
    // 1. 尝试从空闲块中分配
    Block* block = find_free_block(state, size);
    if (block) {
      return allocate_from_block(block, requested_size);
    }
    
    // 2. 尝试垃圾回收
    bool freed_memory = garbage_collect_cached_blocks(state);
    if (freed_memory) {
      block = find_free_block(state, size);
      if (block) {
        return allocate_from_block(block, requested_size);
      }
    }
    
    // 3. 分配新的内存段
    return allocate_new_segment(state, size, requested_size);
  }
  
  // 释放内存（实际加入缓存）
  void deallocate(void* ptr) {
    if (!ptr) return;
    
    std::lock_guard<std::recursive_mutex> lock(mutex_);
    
    auto* block = find_allocated_block(ptr);
    if (!block) {
      throw std::runtime_error("Invalid memory pointer for deallocation");
    }
    
    // 标记为空闲但保留在缓存中
    block->allocated = false;
    block->active = false;
    
    // 将块加入到空闲列表
    auto device = block->device;
    auto& state = device_allocator_states_[device];
    
    if (block->size >= kLargeBuffer) {
      state.large_free_blocks[block->size].insert(block);
    } else {
      state.small_free_blocks[block->size].insert(block);
    }
    
    // 尝试与相邻块合并
    try_merge_blocks(block);
  }
  
 private:
  // 查找合适的空闲块
  Block* find_free_block(DeviceAllocatorState& state, size_t size) {
    // 根据大小选择合适的空闲列表
    auto& free_blocks = (size >= kLargeBuffer) ? 
        state.large_free_blocks : state.small_free_blocks;
    
    // 查找最小的合适块（最佳适配算法）
    auto it = free_blocks.lower_bound(size);
    if (it != free_blocks.end() && !it->second.empty()) {
      Block* block = *it->second.begin();
      it->second.erase(it->second.begin());
      
      // 如果块太大，考虑分割
      if (block->size > size * 2) {
        split_block(block, size);
      }
      
      return block;
    }
    
    return nullptr;
  }
  
  // 垃圾回收算法
  bool garbage_collect_cached_blocks(DeviceAllocatorState& state) {
    size_t freed_bytes = 0;
    
    // 收集过期的空闲块
    auto collect_from_map = [&](auto& free_map) {
      for (auto it = free_map.begin(); it != free_map.end();) {
        auto& block_set = it->second;
        
        for (auto block_it = block_set.begin(); block_it != block_set.end();) {
          auto* block = *block_it;
          
          // 检查块是否可以安全释放
          if (can_release_block(block)) {
            freed_bytes += block->size;
            
            // 实际释放GPU内存
            C10_CUDA_CHECK(cudaFree(block->ptr));
            
            block_it = block_set.erase(block_it);
            delete block;
          } else {
            ++block_it;
          }
        }
        
        if (block_set.empty()) {
          it = free_map.erase(it);
        } else {
          ++it;
        }
      }
    };
    
    // 先收集小块，再收集大块
    collect_from_map(state.small_free_blocks);
    collect_from_map(state.large_free_blocks);
    
    state.reserved_size -= freed_bytes;
    
    return freed_bytes > 0;
  }
  
  // 检查块是否可以释放
  bool can_release_block(Block* block) {
    // 检查是否有活跃的流引用
    for (const auto& stream : block->stream_uses) {
      if (!stream.query()) {
        return false;  // 流还在使用中
      }
    }
    
    // 检查垃圾回收计数器
    return block->gc_counter > 3;  // 经过3次GC周期后可释放
  }
  
  // 内存大小对齐（GPU优化）
  size_t round_size(size_t size) {
    // 小于512字节：对齐到最近的8字节
    if (size < 512) {
      return ((size + 7) / 8) * 8;
    }
    // 大于等于512字节：对齐到最近的512字节
    else {
      return ((size + 511) / 512) * 512;
    }
  }
};

} // namespace c10::cuda::CUDACachingAllocator