Add a mx.metal.device_info (#1060)

* device inof

* add variant

* fix linux

* fix doc

docs/src/python/metal.rst

@@ -7,6 +7,7 @@ Metal
   :toctree: _autosummary
 
   is_available
+  device_info
   get_active_memory
   get_peak_memory
   get_cache_memory

mlx/backend/metal/device.cpp

@@ -556,4 +556,15 @@ void new_stream(Stream stream) {
   }
 }
 
+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());
+  return {
+      {"architecture", arch},
+      {"max_buffer_length", raw_device->maxBufferLength()},
+      {"max_recommended_working_set_size",
+       raw_device->recommendedMaxWorkingSetSize()}};
+}
+
 } // namespace mlx::core::metal

mlx/backend/metal/metal.h

@@ -2,6 +2,8 @@
 
 #pragma once
 
+#include <variant>
+
 #include "mlx/array.h"
 
 namespace mlx::core::metal {
@@ -61,4 +63,8 @@ void clear_cache();
 void start_capture(std::string path = "");
 void stop_capture();
 
+/** Get information about the GPU and system settings. */
+std::unordered_map<std::string, std::variant<std::string, size_t>>
+device_info();
+
 } // namespace mlx::core::metal

mlx/backend/no_metal/metal.cpp

@@ -49,4 +49,10 @@ void start_capture(std::string path) {}
 void stop_capture() {}
 void clear_cache() {}
 
+std::unordered_map<std::string, std::variant<std::string, size_t>>
+device_info() {
+  throw std::runtime_error(
+      "[metal::device_info] Cannot get device info without metal backend");
+};
+
 } // namespace mlx::core::metal

python/src/metal.cpp

@@ -4,6 +4,8 @@
 #include <nanobind/nanobind.h>
 #include <nanobind/stl/optional.h>
 #include <nanobind/stl/string.h>
+#include <nanobind/stl/unordered_map.h>
+#include <nanobind/stl/variant.h>
 
 namespace nb = nanobind;
 using namespace nb::literals;
@@ -116,4 +118,19 @@ void init_metal(nb::module_& m) {
       R"pbdoc(
       Stop a Metal capture.
       )pbdoc");
+  metal.def(
+      "device_info",
+      &metal::device_info,
+      R"pbdoc(
+      Get information about the GPU device and system settings.
+
+      Currently returns:
+
+      * ``architecture``
+      * ``max_buffer_size``
+      * ``max_recommended_working_set_size``
+
+      Returns:
+          dict: A dictionary with string keys and string or integer values.
+      )pbdoc");
 }