Add a mx.metal.device_info (#1060)

* device inof

* add variant

* fix linux

* fix doc

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