[MPS] Add MPSAllocatorInterface to access methods of MPSAllocator (pytorch#94327)

This is a prerequisite for the upcoming PR's for the MPS Modules and Memory Leak Detection features.
Also added pragma once to headers.

Pull Request resolved: pytorch#94327
Approved by: https://github.com/kulinseth

Author: razarmehr

aten/src/ATen/mps/MPSAllocator.h

@@ -1,5 +1,8 @@
 //  Copyright © 2022 Apple Inc.
 
+#pragma once
+
+#include <ATen/mps/MPSAllocatorInterface.h>
 #include <ATen/mps/MPSStream.h>
 #include <cstdio>
 #include <mutex>
@@ -9,27 +12,10 @@
 
 // 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 and remove redundant code.
 namespace at {
 namespace mps {
-
-class IMpsAllocatorCallback {
- public:
-  enum class EventType {
-    ALLOCATED, // buffer got allocated to be used immediately
-    RECYCLED,  // buffer pulled from free list to be reused
-    FREED,     // buffer put to free list for future recycling
-    RELEASED,  // buffer memory released
-  };
-  virtual ~IMpsAllocatorCallback() = default;
-  virtual void executeMPSAllocatorCallback(void* ptr, EventType event) = 0;
-};
-
-// MPS allocator will execute every registered callback when a block of memory is freed.
-C10_DECLARE_REGISTRY(MPSAllocatorCallbacksRegistry, IMpsAllocatorCallback);
-#define REGISTER_MPS_ALLOCATOR_CALLBACK(name, ...) \
-  C10_REGISTER_CLASS(MPSAllocatorCallbacksRegistry, name, __VA_ARGS__);
-
 namespace HeapAllocator {
 
 #define MB(x) round_page(x * 1048576UL)
@@ -263,27 +249,44 @@ class MPSHeapAllocatorImpl
 
   // interface exposed to at::Allocator
   id<MTLBuffer> malloc(size_t size, uint32_t usage);
+  // frees a buffer and returns it into buffer pool
   void free(void* ptr);
+  // releases all the cached buffers and their associated heaps
   void emptyCache();
-  // interface exposed to internal MPS operations
+  // returns true if buffer was allocated from the shared pool
   bool isSharedBuffer(void* ptr);
-  ssize_t getRequestedBufferSize(void* ptr);
+  // get the requested unaligned size of an MTLBuffer
+  ssize_t getUnalignedBufferSize(void* ptr);
+  // set the shape of a base tensor from a view tensor
   void setBufferShape(void* ptr, const IntArrayRef& shape);
+  // retrieve the shape of a base tensor from a view tensor
   IntArrayRef getBufferShape(void* ptr);
+  // allocate a buffer from a specialized pool to import CPU scalars into GPU
   id<MTLBuffer> allocScalarBufferWithValue(void* value, size_t size);
   // 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
   ssize_t getLowWatermarkValue();
-
-  bool getDebugVerbosity() const { return m_debug_verbosity; }
-  size_t getMaxTotalAllowedSize() const { return m_max_total_allowed_size; }
+  // (see m_low_watermark_ratio for description)
+  void setLowWatermarkRatio(double ratio);
+  // (see m_high_watermark_ratio for description)
+  void setHighWatermarkRatio(double ratio);
+  // (see m_low_watermark_limit for description)
   size_t getLowWatermarkLimit() const { return m_low_watermark_limit; }
+  // (see m_max_total_allowed_size for description)
+  size_t getHighWatermarkLimit() const { return m_max_total_allowed_size; }
+  // (see m_total_allocated_memory for description)
+  size_t getTotalAllocatedMemory() const {return m_total_allocated_memory; }
+  // (see enum DebugVerbosity for description)
+  uint32_t getDebugVerbosity() const { return m_debug_verbosity; }
+  // returns the device that we allocate from
   inline id<MTLDevice> Device() const { return m_device; }
 
 private:
   // (see m_high_watermark_ratio for description)
   constexpr static double default_high_watermark_ratio = 1.7;
+  // we set the allowed upper bound to twice the size of recommendedMaxWorkingSetSize.
+  constexpr static double default_high_watermark_upper_bound = 2.0;
   // (see m_low_watermark_ratio for description)
   // on unified memory, we could allocate beyond the recommendedMaxWorkingSetSize
   constexpr static double default_low_watermark_ratio_unified  = 1.4;
@@ -375,17 +378,5 @@ class MPSHeapAllocatorImpl
 };
 
 } // namespace HeapAllocator
-
-// interface exposed to internal MPS operations
-
-// get the requested non-aligned size of an MTL buffer
-ssize_t get_requested_buffer_size(void* ptr);
-// retrieve the shape of a base tensor from a view tensor
-IntArrayRef get_buffer_shape(void* ptr);
-// set the shape of a base tensor from a view tensor
-void set_buffer_shape(void* ptr, const IntArrayRef& shape);
-// allocate a buffer from a specialized pool to import CPU scalars into GPU
-DataPtr allocate_scalar_buffer(void* value, size_t size);
-
 } // namespace mps
 } // namespace at

aten/src/ATen/mps/MPSAllocator.mm

@@ -22,27 +22,35 @@
   static const char *verbosity_str = getenv("PYTORCH_DEBUG_MPS_ALLOCATOR");
   m_debug_verbosity = verbosity_str ? strtol(verbosity_str, nullptr, 0) : DebugVerbosity::SILENT;
 
-  // we set the allowed upper bound to twice the size of recommendedMaxWorkingSetSize.
-  const double high_watermark_upper_bound = 2.0;
-
   static const char *high_watermark_ratio_str = getenv("PYTORCH_MPS_HIGH_WATERMARK_RATIO");
-  m_high_watermark_ratio = high_watermark_ratio_str ? strtod(high_watermark_ratio_str, nullptr) : default_high_watermark_ratio;
-  TORCH_CHECK(m_high_watermark_ratio >= 0.0 && m_high_watermark_ratio <= high_watermark_upper_bound,
-              "invalid high watermark ratio ", m_high_watermark_ratio);
+  const double high_watermark_ratio = high_watermark_ratio_str ? strtod(high_watermark_ratio_str, nullptr) :
+                                                                 default_high_watermark_ratio;
+  setHighWatermarkRatio(high_watermark_ratio);
 
-  m_max_total_allowed_size = (m_high_watermark_ratio == 0.0) ? std::numeric_limits<size_t>::max() :
-                              static_cast<size_t>(m_high_watermark_ratio * (double)max_device_size());
-  // used for comparison with lower_watermark_ratio
-  const double high_watermark_limit = m_high_watermark_ratio == 0.0 ? high_watermark_upper_bound : m_high_watermark_ratio;
   const double default_low_watermark_ratio =  m_device.hasUnifiedMemory ? default_low_watermark_ratio_unified :
                                                                           default_low_watermark_ratio_discrete;
   static const char *low_watermark_ratio_str = getenv("PYTORCH_MPS_LOW_WATERMARK_RATIO");
-  m_low_watermark_ratio = low_watermark_ratio_str ? strtod(low_watermark_ratio_str, nullptr) : default_low_watermark_ratio;
-  TORCH_CHECK(m_low_watermark_ratio >= 0.0 && m_low_watermark_ratio <= high_watermark_limit,
-              "invalid low watermark ratio ", m_low_watermark_ratio);
+  const double low_watermark_ratio = low_watermark_ratio_str ? strtod(low_watermark_ratio_str, nullptr) : default_low_watermark_ratio;
+  setLowWatermarkRatio(low_watermark_ratio);
+}
+
+void MPSHeapAllocatorImpl::setHighWatermarkRatio(double ratio)
+{
+  TORCH_CHECK(ratio >= 0.0 && ratio <= default_high_watermark_upper_bound, "invalid high watermark ratio ", ratio);
+  m_max_total_allowed_size = (ratio == 0.0) ? std::numeric_limits<size_t>::max() :
+                             static_cast<size_t>(ratio * (double)max_device_size());
+  m_high_watermark_ratio = ratio;
+}
+
+void MPSHeapAllocatorImpl::setLowWatermarkRatio(double ratio)
+{
+  // used for comparison with lower_watermark_ratio
+  const double high_watermark_limit = m_high_watermark_ratio == 0.0 ? default_high_watermark_upper_bound : m_high_watermark_ratio;
+  TORCH_CHECK(ratio >= 0.0 && ratio <= high_watermark_limit, "invalid low watermark ratio ", ratio);
   // we use this to detect if there's memory pressure
-  m_low_watermark_limit = (m_low_watermark_ratio == 0.0) ? std::numeric_limits<size_t>::max() :
-                          static_cast<size_t>(m_low_watermark_ratio * (double)max_device_size());
+  m_low_watermark_limit = (ratio == 0.0) ? std::numeric_limits<size_t>::max() :
+                          static_cast<size_t>(ratio * (double)max_device_size());
+  m_low_watermark_ratio = ratio;
 }
 
 HeapBlock* MPSHeapAllocatorImpl::get_free_heap(AllocParams& params)
@@ -470,7 +478,7 @@
   return buffer_block->buffer;
 }
 
-ssize_t MPSHeapAllocatorImpl::getRequestedBufferSize(void* ptr)
+ssize_t MPSHeapAllocatorImpl::getUnalignedBufferSize(void* ptr)
 {
   std::lock_guard<std::mutex> lock(m_mutex);
 
@@ -552,24 +560,24 @@
 }
 
 // MPS allocator struct to be registered with Pytorch
-struct TORCH_API MPSAllocator final : public at::Allocator {
+struct TORCH_API MPSAllocator final : public IMPSAllocator {
 public:
   explicit MPSAllocator(uint32_t Usage) :
       m_has_unified_memory(_getAllocImpl().Device().hasUnifiedMemory), m_usage(Usage)
   {
     if (_getAllocImpl().getDebugVerbosity()) {
       if (!(m_usage & HeapAllocator::UsageFlags::SHARED) || m_has_unified_memory) {
-        const size_t max_total_allowed_size = _getAllocImpl().getMaxTotalAllowedSize();
-        const size_t low_watermark_limit = _getAllocImpl().getLowWatermarkLimit();
+        const size_t high_watermark_limit = _getAllocImpl().getHighWatermarkLimit();
+        const size_t low_watermark_limit  = _getAllocImpl().getLowWatermarkLimit();
         std::cerr << "Initializing "
                   << ((m_usage & HeapAllocator::UsageFlags::SHARED) ? "shared" : "private")
                   << " heap allocator on "
                   << (m_has_unified_memory ? "unified" : "discrete")
                   << " device memory of size "
                   << _getAllocImpl().Device().recommendedMaxWorkingSetSize / 1048576UL << " MB"
                   << " (max allowed: "
-                  << (max_total_allowed_size == std::numeric_limits<size_t>::max() ? "unlimited" :
-                     (to_string(max_total_allowed_size / 1048576UL) + " MB"))
+                  << (high_watermark_limit == std::numeric_limits<size_t>::max() ? "unlimited" :
+                     (to_string(high_watermark_limit / 1048576UL) + " MB"))
                   << ", low watermark: "
                   << (low_watermark_limit == std::numeric_limits<size_t>::max() ? "unlimited" :
                      (to_string(low_watermark_limit / 1048576UL) + " MB"))  << ")\n";
@@ -580,20 +588,28 @@ explicit MPSAllocator(uint32_t Usage) :
   ~MPSAllocator() override {
     _getAllocImpl().emptyCache();
   }
+  DeleterFnPtr raw_deleter() const override { return &Delete; }
 
   DataPtr allocate(const size_t nbytes) const override {
     __block id<MTLBuffer> buf = nbytes > 0 ? _getAllocImpl().malloc(nbytes, m_usage) : nullptr;
     return { buf, buf, &Delete, at::Device(at::DeviceType::MPS, 0)};
   }
-
-  DataPtr allocate_scalar_buffer(void *value, size_t size) const {
+  DataPtr allocScalarBufferWithValue(void *value, size_t size) const override {
     id<MTLBuffer> buf = _getAllocImpl().allocScalarBufferWithValue(value, size);
     return { buf, buf, &Delete, at::Device(at::DeviceType::MPS, 0)};
   }
-
-  DeleterFnPtr raw_deleter() const override { return &Delete; }
-  bool is_shared(void* ptr) const { return _getAllocImpl().isSharedBuffer(ptr); }
-  bool is_shared_storage_supported() const { return m_has_unified_memory; }
+  bool isSharedBuffer(void* ptr) const override { return _getAllocImpl().isSharedBuffer(ptr); }
+  bool isSharedStorageSupported() const override { return m_has_unified_memory; }
+  void emptyCache() const override { _getAllocImpl().emptyCache(); }
+  ssize_t getUnalignedBufferSize(void* ptr) const override { return _getAllocImpl().getUnalignedBufferSize(ptr); }
+  IntArrayRef getBufferShape(void* ptr) const override { return _getAllocImpl().getBufferShape(ptr); }
+  void setBufferShape(void* ptr, const IntArrayRef& shape) const override { _getAllocImpl().setBufferShape(ptr, shape); }
+  size_t getTotalAllocatedMemory() const override { return _getAllocImpl().getTotalAllocatedMemory(); }
+  ssize_t getLowWatermarkValue() const override { return _getAllocImpl().getLowWatermarkValue(); }
+  size_t getLowWatermarkLimit() const override { return _getAllocImpl().getLowWatermarkLimit(); }
+  size_t getHighWatermarkLimit() const override { return _getAllocImpl().getHighWatermarkLimit(); }
+  void setLowWatermarkRatio(double ratio) const override { _getAllocImpl().setLowWatermarkRatio(ratio); }
+  void setHighWatermarkRatio(double ratio) const override { _getAllocImpl().setHighWatermarkRatio(ratio); }
 
 private:
   bool m_has_unified_memory;
@@ -618,41 +634,17 @@ static void Delete(void* ptr) {
 }
 } // anonymous namespace
 
-at::Allocator* getMPSSharedAllocator()
-{
+IMPSAllocator* getIMPSAllocator(bool sharedAllocator) {
+  if (!sharedAllocator) {
+    return &_getPrivateAllocator();
+  }
   auto& sa = _getSharedAllocator();
-  if (sa.is_shared_storage_supported()) {
+  if (sa.isSharedStorageSupported()) {
     return &sa;
   }
-
   return nullptr;
 }
 
-at::Allocator* getMPSPrivateAllocator() {
-  return &_getPrivateAllocator();
-}
-
-// TODO: create MPSHooks interface and move these there.
-ssize_t get_requested_buffer_size(void* ptr) {
-  return _getAllocImpl().getRequestedBufferSize(ptr);
-}
-
-void set_buffer_shape(void* ptr, const IntArrayRef& shape) {
-  _getAllocImpl().setBufferShape(ptr, shape);
-}
-
-IntArrayRef get_buffer_shape(void* ptr) {
-  return _getAllocImpl().getBufferShape(ptr);
-}
-
-DataPtr allocate_scalar_buffer(void *value, size_t size) {
-  return _getPrivateAllocator().allocate_scalar_buffer(value, size);
-}
-
-uint32_t get_adaptive_commit_threshold() {
-  return _getAllocImpl().getLowWatermarkValue();
-}
-
 } // namespace mps
 
 namespace native {
@@ -664,14 +656,14 @@ uint32_t get_adaptive_commit_threshold() {
 bool is_pinned_mps(const Tensor& self, c10::optional<Device> device)
 {
   TORCH_INTERNAL_ASSERT_DEBUG_ONLY(!device.has_value() || device->is_mps());
-  return at::mps::_getSharedAllocator().is_shared(self.storage().data());
+  return at::mps::_getSharedAllocator().isSharedBuffer(self.storage().data());
 }
 
 // torch.pin_memory() implementation
 Tensor _pin_memory_mps(const Tensor& self, c10::optional<Device> device)
 {
   TORCH_INTERNAL_ASSERT_DEBUG_ONLY(!device.has_value() || device->is_mps());
-  auto* shared_allocator = at::mps::getMPSSharedAllocator();
+  auto* shared_allocator = at::mps::getIMPSAllocator(true);
   TORCH_CHECK(shared_allocator, "unable to pin memory on a non-unified memory device");
 
   const size_t storage_size = detail::computeStorageNbytes(self.sizes(), self.strides(), self.dtype().itemsize());

aten/src/ATen/mps/MPSAllocatorInterface.h

@@ -0,0 +1,52 @@
+//  Copyright © 2023 Apple Inc.
+
+#pragma once
+
+#include <c10/core/Allocator.h>
+#include <c10/util/Registry.h>
+#include <ATen/core/ATen_fwd.h>
+
+namespace at {
+namespace mps {
+
+// this is a public interface to access MPSAllocator.
+// Do not declare methods that would depend on MPS or Metal frameworks.
+class IMPSAllocator : public c10::Allocator {
+public:
+  // see the comments in MPSAllocator.h for the description of these methods.
+  virtual void emptyCache() const = 0;
+  virtual ssize_t getUnalignedBufferSize(void* ptr) const = 0;
+  virtual IntArrayRef getBufferShape(void* ptr) const = 0;
+  virtual void setBufferShape(void* ptr, const IntArrayRef& shape) const = 0;
+  virtual bool isSharedBuffer(void* ptr) const = 0;
+  virtual bool isSharedStorageSupported() const = 0;
+  virtual c10::DataPtr allocScalarBufferWithValue(void* value, size_t size) const = 0;
+  virtual void setLowWatermarkRatio(double ratio) const = 0;
+  virtual void setHighWatermarkRatio(double ratio) const = 0;
+  virtual ssize_t getLowWatermarkValue() const = 0;
+  virtual size_t getLowWatermarkLimit() const = 0;
+  virtual size_t getHighWatermarkLimit() const = 0;
+  virtual size_t getTotalAllocatedMemory() const = 0;
+};
+
+class IMpsAllocatorCallback {
+ public:
+  enum class EventType {
+    ALLOCATED, // buffer got allocated to be used immediately
+    RECYCLED,  // buffer pulled from free list to be reused
+    FREED,     // buffer put to free list for future recycling
+    RELEASED,  // buffer memory released
+  };
+  virtual ~IMpsAllocatorCallback() = default;
+  virtual void executeMPSAllocatorCallback(void* ptr, EventType event) = 0;
+};
+
+// MPS allocator will execute every registered callback when a block of memory is freed.
+C10_DECLARE_REGISTRY(MPSAllocatorCallbacksRegistry, IMpsAllocatorCallback);
+#define REGISTER_MPS_ALLOCATOR_CALLBACK(name, ...) \
+  C10_REGISTER_CLASS(MPSAllocatorCallbacksRegistry, name, __VA_ARGS__);
+
+IMPSAllocator* getIMPSAllocator(bool sharedAllocator = false);
+
+} // namespace mps
+} // namespace at

aten/src/ATen/mps/MPSDevice.mm

@@ -3,6 +3,7 @@
 #include <c10/util/CallOnce.h>
 
 #include <ATen/mps/MPSDevice.h>
+#include <ATen/mps/MPSAllocatorInterface.h>
 #include <ATen/mps/IndexKernels.h>
 
 namespace at {
@@ -94,10 +95,8 @@ static inline MTLLanguageVersion getMetalLanguageVersion(const id<MTLDevice>& de
   return _macos13plus;
 }
 
-at::Allocator* getMPSSharedAllocator();
-at::Allocator* getMPSPrivateAllocator();
 at::Allocator* GetMPSAllocator(bool useSharedAllocator) {
-  return useSharedAllocator ? getMPSSharedAllocator() : getMPSPrivateAllocator();
+  return getIMPSAllocator(useSharedAllocator);
 }
 
 bool is_available() {

aten/src/ATen/mps/MPSStream.mm

@@ -1,15 +1,13 @@
 //  Copyright © 2022 Apple Inc.
 
 #include <ATen/mps/MPSStream.h>
+#include <ATen/mps/MPSAllocatorInterface.h>
 
 namespace at {
 namespace mps {
 
 #define USE_COMMIT_AND_CONTINUE 1
 
-// the frequency that we commit the command buffer calculated based on low watermark ratio in MPSAllocator
-uint32_t get_adaptive_commit_threshold();
-
 //-----------------------------------------------------------------
 //  MPSStream
 //-----------------------------------------------------------------
@@ -52,7 +50,7 @@
       break;
     case SyncType::COMMIT_ADAPTIVE:
       // the adaptive commit only commits if we hit the low watermark memory threshold
-      if (get_adaptive_commit_threshold() <= 1) {
+      if (getIMPSAllocator()->getLowWatermarkValue() <= 1) {
 #if USE_COMMIT_AND_CONTINUE
         commitAndContinue();
 #else