Method AllocateRaw

tensorflow/core/framework/tracking_allocator.cc:32–73

      next_allocation_id_(0) {}

void* TrackingAllocator::AllocateRaw(
    size_t alignment, size_t num_bytes,
    const AllocationAttributes& allocation_attr) {
  void* ptr = allocator_->AllocateRaw(alignment, num_bytes, allocation_attr);
  // If memory is exhausted AllocateRaw returns nullptr, and we should
  // pass this through to the caller
  if (nullptr == ptr) {
    return ptr;
  }
  if (allocator_->TracksAllocationSizes()) {
    size_t allocated_bytes = allocator_->AllocatedSize(ptr);
    {
      mutex_lock lock(mu_);
      allocated_ += allocated_bytes;
      high_watermark_ = std::max(high_watermark_, allocated_);
      total_bytes_ += allocated_bytes;
      allocations_.emplace_back(allocated_bytes, Env::Default()->NowMicros());
      ++ref_;
    }
  } else if (track_sizes_locally_) {
    // Call the underlying allocator to try to get the allocated size
    // whenever possible, even when it might be slow. If this fails,
    // use the requested size as an approximation.
    size_t allocated_bytes = allocator_->AllocatedSizeSlow(ptr);
    allocated_bytes = std::max(num_bytes, allocated_bytes);
    mutex_lock lock(mu_);
    next_allocation_id_ += 1;
    Chunk chunk = {num_bytes, allocated_bytes, next_allocation_id_};
    in_use_.emplace(std::make_pair(ptr, chunk));
    allocated_ += allocated_bytes;
    high_watermark_ = std::max(high_watermark_, allocated_);
    total_bytes_ += allocated_bytes;
    allocations_.emplace_back(allocated_bytes, Env::Default()->NowMicros());
    ++ref_;
  } else {
    mutex_lock lock(mu_);
    total_bytes_ += num_bytes;
    allocations_.emplace_back(num_bytes, Env::Default()->NowMicros());
    ++ref_;
  }
  return ptr;
}

size_t TrackingAllocator::BatchAllocateRaw(size_t num, size_t alignment,
    size_t num_bytes, void** ret) {