Pytorch源码阅读-ch3-Tensor&Storage实现

 1
 2
 3
 4
 5
 6
 7
 8
 9
10
11
12
13
14
15
16
17
18
19
20
21
22

struct C10_API TensorImpl : public c10::intrusive_ptr_target 
protected:
  Storage storage_;  
  std::unique_ptr<c10::AutogradMetaInterface> autograd_meta_ = nullptr;  //只有is_variable的时候自动微分才生效
  c10::VariableVersion version_counter_;
  PyObject* pyobj_ = nullptr; // weak reference

  SmallVector<int64_t,5> sizes_;  //每个维度的size
  SmallVector<int64_t,5> strides_; //每个维度元素的间隔

  int64_t storage_offset_ = 0;  
  int64_t numel_ = 1;  //元素总个数
  caffe2::TypeMeta data_type_;
  c10::optional<c10::Device> device_opt_;

  TensorTypeId type_id_;
  bool is_contiguous_ = true;  //在内存中是否连续
  bool is_wrapped_number_ = false;

  bool allow_tensor_metadata_change_ = true; //拒绝 detach产出的tensor被更改
  bool reserved_ = false;
};

1
2

  Tensor a:  00 01 02 03 04 05 10 11 12 13 14 15 20 21 22 23 24 25 
  Tensor b:  00 01 02 03  x  x 10 11 12 13  x  x 20 21 22 23  x  x

1
2
3
4
5
6
7
8

struct C10_API StorageImpl final : public c10::intrusive_ptr_target {
 private:
  caffe2::TypeMeta data_type_;  // 数据类型
  DataPtr data_ptr_;            // 指向存储数据的内存块
  int64_t numel_;               // 数据总数
  bool resizable_;
  Allocator* allocator_;        // 内存分配器
}

 1
 2
 3
 4
 5
 6
 7
 8
 9
10
11
12
13
14
15
16
17
18
19
20

struct C10_API Allocator {
  virtual ~Allocator() = default;

  virtual DataPtr allocate(size_t n) const = 0;

  //using DeleterFnPtr = void (*)(void*);
  virtual DeleterFnPtr raw_deleter() const { //这里需要返回一个安全的资源释放函数
    return nullptr;
  }
  void* raw_allocate(size_t n) {
    auto dptr = allocate(n);
    AT_ASSERT(dptr.get() == dptr.get_context());
    return dptr.release_context();
  }
  void raw_deallocate(void* ptr) {
    auto d = raw_deleter();
    AT_ASSERT(d);
    d(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

struct C10_API DefaultCPUAllocator final : at::Allocator {
  DefaultCPUAllocator() {}
  ~DefaultCPUAllocator() override {}
  at::DataPtr allocate(size_t nbytes) const override {  //virtual override
    void* data = alloc_cpu(nbytes);  //类外函数，为了兼容不同平台的alloc函数
    if (FLAGS_caffe2_report_cpu_memory_usage && nbytes > 0) {
      getMemoryAllocationReporter().New(data, nbytes); //使用unorderd_map 
      return {data, data, &ReportAndDelete, at::Device(at::DeviceType::CPU)};
    }
    return {data, data, &free_cpu, at::Device(at::DeviceType::CPU)};
  }

  static void ReportAndDelete(void* ptr) {
    if (!ptr) {
      return;
    }
    getMemoryAllocationReporter().Delete(ptr);
    free_cpu(ptr);
  }

  at::DeleterFnPtr raw_deleter() const override { //virtual override
    if (FLAGS_caffe2_report_cpu_memory_usage) {
      return &ReportAndDelete;  
    }
    return &free_cpu;
  }

 protected:
  static MemoryAllocationReporter& getMemoryAllocationReporter() {  //这个结构记录了所有已分配的指针与对应的size
    static MemoryAllocationReporter reporter_;  //静态，使用的时候加了互斥锁
    return reporter_;
  }
};