HostBlockPool Class — pytorch Architecture

Source Code

aten/src/ATen/core/CachingHostAllocator.h lines 229–246

template <typename S_, typename E_, typename B_>
struct HostBlockPool {
  HostBlockPool() = default;
  HostBlockPool(const HostBlockPool&) = delete;
  HostBlockPool& operator=(const HostBlockPool&) = delete;

  alignas(hardware_destructive_interference_size) std::mutex blocks_mutex_;
  ska::flat_hash_set<B_*> blocks_; // all blocks in this pool
  ska::flat_hash_map<void*, B_*> ptr_to_block_;

  // Per-size free lists guarded by their own mutexes.
  alignas(hardware_destructive_interference_size) std::vector<FreeBlockList<B_>> free_list_ =
      std::vector<FreeBlockList<B_>>(MAX_SIZE_INDEX);

  // Events pending for blocks in this pool.
  alignas(hardware_destructive_interference_size) std::mutex events_mutex_;
  std::deque<std::pair<E_, B_*>> events_;
};