HostBlock Class — pytorch Architecture

Source Code

aten/src/ATen/core/CachingHostAllocator.h lines 28–43

template <typename S>
struct HostBlock {
  // constructor for search key
  HostBlock(size_t size) : size_(size) {}

  HostBlock(size_t size, void* ptr) : size_(size), ptr_(ptr) {}

  std::mutex mutex_;
  size_t size_{0}; // block size in bytes
  void* ptr_{nullptr}; // memory address
  bool allocated_{false}; // in-use flag
  size_t event_count_{0}; // number of related events
  ska::flat_hash_set<S> streams_; // streams on which the block was used
  c10::MempoolId_t owning_pool_{0,0}; // never changes after construction, so we don't need a mutex to guard this
  bool was_allocated_during_stream_capture_;
};