GPU discovery (#3055)

Co-authored-by: Angelos Katharopoulos <[email protected]>

python/src/device.cpp

@@ -1,9 +1,11 @@
-// Copyright © 2023-2024 Apple Inc.
+// Copyright © 2023-2025 Apple Inc.
 
 #include <sstream>
 
 #include <nanobind/nanobind.h>
 #include <nanobind/stl/string.h>
+#include <nanobind/stl/unordered_map.h>
+#include <nanobind/stl/variant.h>
 
 #include "mlx/device.h"
 #include "mlx/utils.h"
@@ -63,4 +65,34 @@ void init_device(nb::module_& m) {
       &mx::is_available,
       "device"_a,
       R"pbdoc(Check if a back-end is available for the given device.)pbdoc");
+  m.def(
+      "device_count",
+      &mx::device_count,
+      "device_type"_a,
+      R"pbdoc(
+      Get the number of available devices for the given device type.
+
+      Args:
+          device_type (DeviceType): The type of device to query (cpu or gpu).
+
+      Returns:
+          int: Number of devices.
+      )pbdoc");
+  m.def(
+      "device_info",
+      &mx::device_info,
+      nb::arg("d") = mx::default_device(),
+      R"pbdoc(
+      Get information about a device.
+
+      Returns a dictionary with device properties. Available keys depend
+      on the backend and device type. Common keys include ``device_name``,
+      ``architecture``, and ``total_memory`` (or ``memory_size``).
+
+      Args:
+          d (Device): The device to query (defaults to the default device).
+
+      Returns:
+          dict: Device information.
+      )pbdoc");
 }

python/src/metal.cpp

@@ -9,6 +9,7 @@
 #include <nanobind/stl/vector.h>
 
 #include "mlx/backend/metal/metal.h"
+#include "mlx/device.h"
 #include "mlx/memory.h"
 #include "python/src/small_vector.h"
 
@@ -90,21 +91,8 @@ void init_metal(nb::module_& m) {
       R"pbdoc(
       Stop a Metal capture.
       )pbdoc");
-  metal.def(
-      "device_info",
-      &mx::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``
-      * ``memory_size``
-      * ``resource_limit``
-
-      Returns:
-          dict: A dictionary with string keys and string or integer values.
-      )pbdoc");
+  metal.def("device_info", []() {
+    DEPRECATE("mx.metal.device_info", "mx.device_info");
+    return mx::device_info(mx::Device(mx::Device::gpu, 0));
+  });
 }

python/tests/test_device.py

@@ -113,5 +113,38 @@ class TestStream(mlx_tests.MLXTestCase):
         self.assertEqual(a.item(), b.item())
 
 
+class TestDeviceInfo(mlx_tests.MLXTestCase):
+    def test_device_count(self):
+        cpu_count = mx.device_count(mx.cpu)
+        self.assertIsInstance(cpu_count, int)
+        self.assertEqual(cpu_count, 1)
+
+        gpu_count = mx.device_count(mx.gpu)
+        self.assertIsInstance(gpu_count, int)
+        self.assertGreaterEqual(gpu_count, 0)
+
+    def test_device_info_cpu(self):
+        info = mx.device_info(mx.cpu)
+        self.assertIsInstance(info, dict)
+        self.assertIn("device_name", info)
+        self.assertTrue(len(info["device_name"]) > 0)
+        self.assertIn("architecture", info)
+
+    @unittest.skipIf(not mx.is_available(mx.gpu), "GPU is not available")
+    def test_device_info_gpu(self):
+        gpu_count = mx.device_count(mx.gpu)
+        for i in range(gpu_count):
+            info = mx.device_info(mx.Device(mx.gpu, i))
+            self.assertIsInstance(info, dict)
+            self.assertIn("device_name", info)
+            self.assertTrue(len(info["device_name"]) > 0)
+            self.assertIn("architecture", info)
+
+    def test_device_info_default(self):
+        info = mx.device_info()
+        self.assertIsInstance(info, dict)
+        self.assertIn("device_name", info)
+
+
 if __name__ == "__main__":
     mlx_tests.MLXTestRunner()

python/tests/test_memory.py

@@ -54,7 +54,7 @@ class TestMemory(mlx_tests.MLXTestCase):
         old_limit = mx.set_wired_limit(0)
         self.assertEqual(old_limit, 1000)
 
-        max_size = mx.metal.device_info()["max_recommended_working_set_size"]
+        max_size = mx.device_info(mx.gpu)["max_recommended_working_set_size"]
         with self.assertRaises(ValueError):
             mx.set_wired_limit(max_size + 10)