change initial memory limits and add memory size to device info (#1064)

mlx/backend/metal/allocator.cpp

@@ -140,10 +140,15 @@ void BufferCache::remove_from_list(BufferCache::BufferHolder* to_remove) {
 
 MetalAllocator::MetalAllocator()
     : device_(device(mlx::core::Device::gpu).mtl_device()),
-      buffer_cache_(device_),
-      block_limit_(1.5 * device_->recommendedMaxWorkingSetSize()),
-      gc_limit_(0.95 * device_->recommendedMaxWorkingSetSize()),
-      max_pool_size_(block_limit_) {}
+      buffer_cache_(device_) {
+  auto memsize = std::get<size_t>(device_info()["memory_size"]);
+  block_limit_ =
+      std::min(1.5 * device_->recommendedMaxWorkingSetSize(), 0.95 * memsize);
+  gc_limit_ = std::min(
+      static_cast<size_t>(0.95 * device_->recommendedMaxWorkingSetSize()),
+      block_limit_);
+  max_pool_size_ = block_limit_;
+}
 
 size_t MetalAllocator::set_cache_limit(size_t limit) {
   std::swap(limit, max_pool_size_);

mlx/backend/metal/device.cpp

@@ -5,6 +5,8 @@
 #include <filesystem>
 #include <sstream>
 
+#include <sys/sysctl.h>
+
 #define NS_PRIVATE_IMPLEMENTATION
 #define CA_PRIVATE_IMPLEMENTATION
 #define MTL_PRIVATE_IMPLEMENTATION
@@ -560,11 +562,19 @@ std::unordered_map<std::string, std::variant<std::string, size_t>>
 device_info() {
   auto raw_device = device(default_device()).mtl_device();
   auto arch = std::string(raw_device->architecture()->name()->utf8String());
+
+  int mib[] = {CTL_HW, HW_MEMSIZE};
+  size_t memsize = 0;
+  size_t length = sizeof(memsize);
+
+  sysctl(mib, 2, &memsize, &length, NULL, 0);
+
   return {
       {"architecture", arch},
       {"max_buffer_length", raw_device->maxBufferLength()},
       {"max_recommended_working_set_size",
-       raw_device->recommendedMaxWorkingSetSize()}};
+       raw_device->recommendedMaxWorkingSetSize()},
+      {"memory_size", memsize}};
 }
 
 } // namespace mlx::core::metal

python/src/metal.cpp

@@ -129,6 +129,7 @@ void init_metal(nb::module_& m) {
       * ``architecture``
       * ``max_buffer_size``
       * ``max_recommended_working_set_size``
+      * ``memory_size``
 
       Returns:
           dict: A dictionary with string keys and string or integer values.