JAX 是怎么分配内存的

// 一个实现了 'best-fit with coalescing'（最佳适配 + 合并）算法的内存分配器。
// 本质上是 Doug Lea 的 malloc（dlmalloc）的一个极简版本。
//
// 这个分配器的目标是通过合并（coalescing）来支持碎片整理。
// 我们的一个前提假设是：使用它的进程几乎拥有全部内存，
// 并且几乎所有的内存分配请求都通过这个接口进行。
class BFCAllocator : public Allocator {

class AllocationRegion {
  // ...
  void* ptr_ = nullptr;       // region起始地址
  size_t memory_size_ = 0;    // region大小
  void* end_ptr_ = nullptr;   // region结束地址

  // 这个数组的大小是memory_size_ / 256B，代表每256B一个槽位，每个槽位存储一个ChunkHandle
  // 但只有Chunk起始地址对应的槽位存着有效的ChunkHandle（用户拿到的ptr一定是某个Chunk的起始地址）。
  // 这样就能通过ptr以O(1)反查出ChunkHandle，进而在chunks_数组中找到对应的Chunk
  std::vector<ChunkHandle> handles_; 
};

static constexpr size_t kMinAllocationBits = 8;
static constexpr size_t kMinAllocationSize = 1 << kMinAllocationBits;  // 256B

// 一个 Chunk 指向一块内存，这块内存要么整块空闲，要么整块被某一次用户分配占用。
//
// 所有 Chunk 串成一个双向链表，
// prev/next 指针指向物理上紧邻的前后 Chunk。
struct Chunk {
  size_t size = 0;            // 这块内存的实际大小
  size_t requested_size = 0;  // 用户真正请求的大小（size >= requested_size）
  int64_t allocation_id = -1; // -1表示free，>0表示in use
  void* ptr = nullptr;        // 内存起始地址

  ChunkHandle prev = kInvalidChunkHandle;  // 物理上紧邻的前一个Chunk
  ChunkHandle next = kInvalidChunkHandle;  // 物理上紧邻的后一个Chunk
  BinNum bin_num = kInvalidBinNum;         // free时所在的Bin编号

  bool in_use() const { return allocation_id != -1; }
};

// 一个 Bin 是一组大小相近的空闲 Chunk 的集合。
// 已分配的 Chunk 永远不会出现在 Bin 里。
struct Bin {
  // 这个 Bin 里的所有 Chunk 大小都 >= bin_size。
  size_t bin_size = 0;

  // Bin 内的空闲 Chunk 列表，按 Chunk 大小排序。
  FreeChunkSet free_chunks;  // absl::btree_set<ChunkHandle, ChunkComparator>
};

class ChunkComparator {
   // 先按 size 排序，size 相同时用指针地址作为决胜判据（tie breaker）。
   bool operator()(const ChunkHandle ha, const ChunkHandle hb) const
         ABSL_NO_THREAD_SAFETY_ANALYSIS {
      const Chunk* a = allocator_->ChunkFromHandle(ha);
      const Chunk* b = allocator_->ChunkFromHandle(hb);
      if (a->size != b->size) {
         return a->size < b->size;
      }
      return a->ptr < b->ptr;
   }

void* BFCAllocator::AllocateRawInternal(size_t alignment, size_t num_bytes, ...) {
  size_t rounded_bytes = RoundedBytes(num_bytes);  // 向上取整到256B的倍数
  BinNum bin_num = BinNumForSize(rounded_bytes);    // 落到哪个Bin

  absl::MutexLock l(mutex_);
  // 1. 先在现有的空闲Chunk里找一块够用的
  void* ptr = FindChunkPtr(bin_num, rounded_bytes, num_bytes, alignment, ...);
  if (ptr != nullptr) return ptr;

  // 2. 找不到 → Extend()向底层要一块大的，再找一次
  if (Extend(alignment, rounded_bytes)) {
    ptr = FindChunkPtr(bin_num, rounded_bytes, num_bytes, alignment, ...);
    if (ptr != nullptr) return ptr;
  }

  // 3. 还是不行 → 回收空闲region，最终仍失败就是OOM
  // ...
}

void* BFCAllocator::FindChunkPtr(BinNum bin_num, size_t rounded_bytes,
                                 size_t num_bytes, size_t alignment, ...) {
  // 从目标Bin开始，向更大的Bin方向扫描
  for (; bin_num < kNumBins; bin_num++) {
    Bin* b = BinFromIndex(bin_num);
    // Bin内是按size升序的有序set，第一个装得下的就是这个Bin里最小够用的
    for (auto citer = b->free_chunks.begin(); citer != b->free_chunks.end();
         ++citer) {
      ChunkHandle h = (*citer);
      Chunk* chunk = ChunkFromHandle(h);

      if (chunk->size >= rounded_bytes /* + 对齐padding */) {
        // 先把它从空闲Bin里摘掉
        RemoveFreeChunkIterFromBin(&b->free_chunks, citer);

        // 如果这块远大于所需，就Split（见下面的SplitChunk）
        if (chunk->size >= rounded_bytes * 2 || /* 剩余超过128MB */) {
          SplitChunk(h, rounded_bytes);
          chunk = ChunkFromHandle(h);
        }

        chunk->requested_size = num_bytes;
        chunk->allocation_id = next_allocation_id_++;  // 标记为in use
        return chunk->ptr;
      }
    }
  }
  return nullptr;  // 所有Bin都找不到 → 交回上层去Extend()
}

bool BFCAllocator::Extend(size_t alignment, size_t rounded_bytes) {
  size_t available_bytes = memory_limit_ - *stats_.pool_bytes;
  available_bytes = (available_bytes / kMinAllocationSize) * kMinAllocationSize;
  if (rounded_bytes > available_bytes) {
    return false;  // 预算上限不够的话(比如硬件上限)，直接失败
  }

  // 不够就按2的幂往上翻，直到能装下这次请求
  bool increased_allocation = false;
  while (rounded_bytes > curr_region_allocation_bytes_) {
    curr_region_allocation_bytes_ *= 2;
    increased_allocation = true;
  }

  size_t bytes = std::min(curr_region_allocation_bytes_, available_bytes);
  size_t bytes_received;
  void* mem_addr = sub_allocator_->Alloc(alignment, bytes, &bytes_received);
  // ... 如果Alloc失败，按0.9的系数不断回退重试 ...

  // 把整块新内存先包成一个大的 Chunk，后续再按需切分。
  ChunkHandle h = AllocateChunk();
  BFCAllocator::Chunk* c = ChunkFromHandle(h);
  c->ptr = mem_addr;
  c->size = bytes_received;
  c->allocation_id = -1;          // -1 表示空闲
  c->prev = kInvalidChunkHandle;
  c->next = kInvalidChunkHandle;

  region_manager_.set_handle(c->ptr, h);
  // 可能和相邻region合并，并把这个大Chunk插进对应Bin
  InsertFreeChunkIntoBin(TryToCoalesce(h, /*ignore_freed_at=*/false));

void BFCAllocator::SplitChunk(BFCAllocator::ChunkHandle h, size_t num_bytes) {
  ChunkHandle h_new_chunk = AllocateChunk();
  Chunk* c = ChunkFromHandle(h);
  CHECK(!c->in_use() && (c->bin_num == kInvalidBinNum));

  // 新Chunk从 c->ptr + num_bytes 处开始
  BFCAllocator::Chunk* new_chunk = ChunkFromHandle(h_new_chunk);
  new_chunk->ptr = static_cast<void*>(static_cast<char*>(c->ptr) + num_bytes);
  region_manager_.set_handle(new_chunk->ptr, h_new_chunk);

  new_chunk->size = c->size - num_bytes;  // 剩余部分
  c->size = num_bytes;                     // 前半段缩小到请求大小
  new_chunk->allocation_id = -1;

  // 维护物理相邻链表： c <-> neighbor  变成  c <-> new_chunk <-> neighbor
  BFCAllocator::ChunkHandle h_neighbor = c->next;
  new_chunk->prev = h;
  new_chunk->next = h_neighbor;
  c->next = h_new_chunk;
  if (h_neighbor != kInvalidChunkHandle) {
    Chunk* c_neighbor = ChunkFromHandle(h_neighbor);
    c_neighbor->prev = h_new_chunk;
  }

  InsertFreeChunkIntoBin(h_new_chunk);  // 剩余部分作为新空闲Chunk入Bin
}

const int64_t max_internal_fragmentation_bytes =
    (opts_.fragmentation_fraction > 0.0)
        ? opts_.fragmentation_fraction * memory_limit_
        : 128 << 20;  // 默认128MB

if (chunk->size >= rounded_bytes * 2 ||
    static_cast<int64_t>(chunk->size) - rounded_bytes >=
        max_internal_fragmentation_bytes) {
  SplitChunk(h, rounded_bytes);
}

void BFCAllocator::Merge(BFCAllocator::ChunkHandle h1,
                         BFCAllocator::ChunkHandle h2) {
  Chunk* c1 = ChunkFromHandle(h1);
  Chunk* c2 = ChunkFromHandle(h2);
  CHECK(!c1->in_use() && !c2->in_use());  // 两块都必须空闲

  // 处理邻居链表: c1 <-> c2 <-> c3  合并成  c1 <-> c3
  BFCAllocator::ChunkHandle h3 = c2->next;
  c1->next = h3;
  if (h3 != kInvalidChunkHandle) {
    ChunkFromHandle(h3)->prev = h1;
  }

  c1->size += c2->size;  // c1 吞并 c2 的大小
  DeleteChunk(h2);       // c2 的metadata被回收
}

BFCAllocator::ChunkHandle BFCAllocator::TryToCoalesce(ChunkHandle h, ...) {
  Chunk* c = ChunkFromHandle(h);
  ChunkHandle coalesced_chunk = h;

  // 右邻居空闲 → 合并进来
  if (c->next != kInvalidChunkHandle && !ChunkFromHandle(c->next)->in_use()) {
    RemoveFreeChunkFromBin(c->next);
    Merge(h, c->next);
  }
  // 前邻居空闲 → 把自己合并进去
  if (c->prev != kInvalidChunkHandle && !ChunkFromHandle(c->prev)->in_use()) {
    coalesced_chunk = c->prev;
    RemoveFreeChunkFromBin(c->prev);
    Merge(c->prev, h);
  }
  return coalesced_chunk;
}