Method AllocateRaw

tensorflow/core/framework/experimental_pmem_allocator.cc:209–277

}

void* ExperimentalPMemAllocator::AllocateRaw(size_t alignment, size_t size) {
  void* ret = nullptr;
  int t_id = MaybeInitAccessThread();
  if (t_id < 0) {
    LOG(FATAL) << "Experimental PMem Allocator: Too many thread access "
                  "allocator! max threads: "
               << kMaxAccessThreads;
    return nullptr;
  }
  uint32_t b_size = Size2BlockSize(size);
  uint32_t aligned_size = b_size * block_size_;
  if (aligned_size > max_allocation_size_ || aligned_size == 0) {
    LOG(FATAL) << "Experimental PMem Allocator: Allocating size: " << size
               << ", is 0 or larger than PMem allocator max allocation size "
               << max_allocation_size_;
    return nullptr;
  }
  auto& thread_cache = thread_cache_[t_id];
  for (auto i = b_size; i < thread_cache.freelists.size(); i++) {
    if (thread_cache.segments[i].size < aligned_size) {
      // Fetch free list from pool
      {
        std::lock_guard<spin_lock> lg(thread_cache.locks[i]);
        if (thread_cache.freelists[i].empty()) {
          pool_.FetchEntryList(thread_cache.freelists[i], i);
        }
        // Get space from free list
        if (thread_cache.freelists[i].size() > 0) {
          ret = thread_cache.freelists[i].back();
          thread_cache.freelists[i].pop_back();
          break;
        }
      }
      // Allocate a new segment for requesting block size
      if (!AllocateSegmentSpace(&thread_cache.segments[b_size], b_size)) {
        continue;
      } else {
        i = b_size;
      }
    }
    assert(thread_cache.segments[i].size >= aligned_size);
    ret = thread_cache.segments[i].addr;
    thread_cache.segments[i].size -= aligned_size;
    thread_cache.segments[i].addr =
        (char*)thread_cache.segments[i].addr + aligned_size;
    break;
  }

  if (pmem_allocator_collect_stats) {
    const std::size_t alloc_size = AllocatedSize(ret);
    mutex_lock l(mu_);
    ++stats_.num_allocs;
    stats_.bytes_in_use += alloc_size;
    stats_.peak_bytes_in_use =
        std::max<int64>(stats_.peak_bytes_in_use, stats_.bytes_in_use);
    stats_.largest_alloc_size =
        std::max<int64>(stats_.largest_alloc_size, alloc_size);