PyTorch内存管理：Caching Allocator（二）一文读懂Block
前言 PyTorch内存管理：Caching Allocator 的第一部分是对PyTorch内存管理的基本情况的概括。这篇文章讲从代码入手，尝试理解PyTorch Caching Allocator 中最核心的数据结构：Block，以及它在实际内存分配中是如何被管理的。 Block 先通过代码看看它的定义： struct Block { c10::DeviceIndex device; // gpu cudaStream_t stream; // allocation stream stream_set stream_uses; // streams on which the block was used size_t size; // block size in bytes size_t requested_size; // memory originally requested BlockPool* pool{nullptr}; // owning memory pool void* ptr{nullptr}; // memory address bool allocated{false}; // in-use flag bool mapped{true}; // is the virtual address range this Block references // backed by physical pages. Always true when // expandable_segment_ is null. When false // This Block will be aligned to the segment size // of its expandable_segment_. Block* prev{nullptr}; // prev block if split from a larger allocation Block* next{nullptr}; // next block if split from a larger allocation int event_count{0}; // number of outstanding CUDA events int64_t gc_count_base{0}; // get_free_blocks_call_count when Block is inserted std::shared_ptr<GatheredContext> context_when_allocated; // only set for the first block in the segment (when prev == null) // this records the frame information when cudaMalloc was called // whereas context_when_allocated records the last time we handed this // memory out from our cache. std::shared_ptr<GatheredContext> context_when_segment_allocated; ExpandableSegment* expandable_segment_{nullptr}; device、size、ptr：这几个参数比较好理解，由于是内存管理，就一定需要涉及到分配的设备、分配的内存大小以及在设备上的地址； requested_size：分配时请求的实际大小； stream：代表分配时的stream，默认所有操作会在该流进行； stream_uses：使用多流时记录使用过的流； prev、next：说明Block本质是一个链表，这也是为什么可以很方便的在Block内部执行切分、重组操作； event_count：未完成的cuda_event计数； 还剩一些涉及到虚拟内存、可扩展段以及资源管理的描述，我们放到后文去深究。总之，这几个属性就已经决定了如何申请、释放一个block，以及存在异步冲突怎么解决。 ...