From a828e769be4566d932e1f03d25370ab785f5e337 Mon Sep 17 00:00:00 2001 From: Daniel Hiltgen Date: Mon, 26 Jan 2026 09:54:13 -0800 Subject: [PATCH] GPU discovery (#3055) Co-authored-by: Angelos Katharopoulos --- docs/src/python/devices_and_streams.rst | 2 + mlx/backend/cpu/CMakeLists.txt | 2 +- mlx/backend/cpu/available.cpp | 11 -- mlx/backend/cpu/available.h | 9 - mlx/backend/cpu/device_info.cpp | 59 ++++++ mlx/backend/cpu/device_info.h | 28 +++ mlx/backend/cuda/CMakeLists.txt | 2 +- mlx/backend/cuda/cuda.cpp | 11 -- mlx/backend/cuda/cuda.h | 9 + mlx/backend/cuda/device_info.cpp | 232 ++++++++++++++++++++++++ mlx/backend/cuda/eval.cpp | 5 - mlx/backend/gpu/available.h | 11 -- mlx/backend/gpu/device_info.h | 36 ++++ mlx/backend/metal/CMakeLists.txt | 1 + mlx/backend/metal/allocator.cpp | 14 +- mlx/backend/metal/device_info.cpp | 58 ++++++ mlx/backend/metal/eval.cpp | 5 - mlx/backend/metal/metal.cpp | 34 ---- mlx/backend/no_cpu/CMakeLists.txt | 2 +- mlx/backend/no_cpu/available.cpp | 11 -- mlx/backend/no_cpu/device_info.cpp | 22 +++ mlx/backend/no_gpu/CMakeLists.txt | 1 + mlx/backend/no_gpu/device_info.cpp | 22 +++ mlx/backend/no_gpu/eval.cpp | 6 +- mlx/device.cpp | 31 +++- mlx/device.h | 25 ++- mlx/mlx.h | 2 +- mlx/scheduler.cpp | 2 +- python/src/device.cpp | 34 +++- python/src/metal.cpp | 22 +-- python/tests/test_device.py | 33 ++++ python/tests/test_memory.py | 2 +- 32 files changed, 608 insertions(+), 136 deletions(-) delete mode 100644 mlx/backend/cpu/available.cpp delete mode 100644 mlx/backend/cpu/available.h create mode 100644 mlx/backend/cpu/device_info.cpp create mode 100644 mlx/backend/cpu/device_info.h delete mode 100644 mlx/backend/cuda/cuda.cpp create mode 100644 mlx/backend/cuda/device_info.cpp delete mode 100644 mlx/backend/gpu/available.h create mode 100644 mlx/backend/gpu/device_info.h create mode 100644 mlx/backend/metal/device_info.cpp delete mode 100644 mlx/backend/no_cpu/available.cpp create mode 100644 mlx/backend/no_cpu/device_info.cpp create mode 100644 mlx/backend/no_gpu/device_info.cpp diff --git a/docs/src/python/devices_and_streams.rst b/docs/src/python/devices_and_streams.rst index 2a5adc05..78dd3e56 100644 --- a/docs/src/python/devices_and_streams.rst +++ b/docs/src/python/devices_and_streams.rst @@ -17,3 +17,5 @@ Devices and Streams set_default_stream stream synchronize + device_count + device_info diff --git a/mlx/backend/cpu/CMakeLists.txt b/mlx/backend/cpu/CMakeLists.txt index 9d322c4c..f2c96949 100644 --- a/mlx/backend/cpu/CMakeLists.txt +++ b/mlx/backend/cpu/CMakeLists.txt @@ -40,7 +40,7 @@ add_dependencies(mlx cpu_compiled_preamble) target_sources( mlx - PRIVATE ${CMAKE_CURRENT_SOURCE_DIR}/available.cpp + PRIVATE ${CMAKE_CURRENT_SOURCE_DIR}/device_info.cpp ${CMAKE_CURRENT_SOURCE_DIR}/arg_reduce.cpp ${CMAKE_CURRENT_SOURCE_DIR}/binary.cpp ${CMAKE_CURRENT_SOURCE_DIR}/conv.cpp diff --git a/mlx/backend/cpu/available.cpp b/mlx/backend/cpu/available.cpp deleted file mode 100644 index 0449d49b..00000000 --- a/mlx/backend/cpu/available.cpp +++ /dev/null @@ -1,11 +0,0 @@ -// Copyright © 2025 Apple Inc. - -#include "mlx/backend/cpu/available.h" - -namespace mlx::core::cpu { - -bool is_available() { - return true; -} - -} // namespace mlx::core::cpu diff --git a/mlx/backend/cpu/available.h b/mlx/backend/cpu/available.h deleted file mode 100644 index 1df95def..00000000 --- a/mlx/backend/cpu/available.h +++ /dev/null @@ -1,9 +0,0 @@ -// Copyright © 2025 Apple Inc. - -#pragma once - -namespace mlx::core::cpu { - -bool is_available(); - -} // namespace mlx::core::cpu diff --git a/mlx/backend/cpu/device_info.cpp b/mlx/backend/cpu/device_info.cpp new file mode 100644 index 00000000..59da8664 --- /dev/null +++ b/mlx/backend/cpu/device_info.cpp @@ -0,0 +1,59 @@ +// Copyright © 2026 Apple Inc. + +#include "mlx/backend/cpu/device_info.h" + +#ifdef __APPLE__ +#include +#endif + +namespace mlx::core::cpu { + +namespace { + +// Get CPU architecture string +std::string get_cpu_architecture() { +#if defined(__aarch64__) || defined(__arm64__) + return "arm64"; +#elif defined(__x86_64__) || defined(_M_X64) + return "x86_64"; +#elif defined(__i386__) || defined(__i386) || defined(_M_IX86) + return "x86"; +#elif defined(__arm__) || defined(_M_ARM) + return "arm"; +#else + return "unknown"; +#endif +} + +// Get CPU device name +std::string get_cpu_name() { +#ifdef __APPLE__ + char model[256]; + size_t len = sizeof(model); + if (sysctlbyname("machdep.cpu.brand_string", &model, &len, NULL, 0) == 0) { + return std::string(model); + } +#endif + return get_cpu_architecture(); +} + +} // anonymous namespace + +bool is_available() { + return true; +} + +int device_count() { + return 1; +} + +const std::unordered_map>& +device_info(int /* device_index */) { + static auto info = + std::unordered_map>{ + {"device_name", get_cpu_name()}, + {"architecture", get_cpu_architecture()}}; + return info; +} + +} // namespace mlx::core::cpu diff --git a/mlx/backend/cpu/device_info.h b/mlx/backend/cpu/device_info.h new file mode 100644 index 00000000..1e232334 --- /dev/null +++ b/mlx/backend/cpu/device_info.h @@ -0,0 +1,28 @@ +// Copyright © 2026 Apple Inc. + +#pragma once + +#include +#include +#include + +namespace mlx::core::cpu { + +bool is_available(); + +/** + * Get the number of available CPU devices. + * + * For CPU, always returns 1. + */ +int device_count(); + +/** + * Get CPU device information. + * + * Returns a map with basic CPU device properties. + */ +const std::unordered_map>& +device_info(int device_index = 0); + +} // namespace mlx::core::cpu diff --git a/mlx/backend/cuda/CMakeLists.txt b/mlx/backend/cuda/CMakeLists.txt index 7f9b6c86..5fb5096e 100644 --- a/mlx/backend/cuda/CMakeLists.txt +++ b/mlx/backend/cuda/CMakeLists.txt @@ -19,8 +19,8 @@ target_sources( ${CMAKE_CURRENT_SOURCE_DIR}/conv/gemm_conv.cu ${CMAKE_CURRENT_SOURCE_DIR}/conv/gemm_grouped_conv.cu ${CMAKE_CURRENT_SOURCE_DIR}/cublas_utils.cpp - ${CMAKE_CURRENT_SOURCE_DIR}/cuda.cpp ${CMAKE_CURRENT_SOURCE_DIR}/cudnn_utils.cpp + ${CMAKE_CURRENT_SOURCE_DIR}/device_info.cpp ${CMAKE_CURRENT_SOURCE_DIR}/custom_kernel.cpp ${CMAKE_CURRENT_SOURCE_DIR}/device.cpp ${CMAKE_CURRENT_SOURCE_DIR}/distributed.cu diff --git a/mlx/backend/cuda/cuda.cpp b/mlx/backend/cuda/cuda.cpp deleted file mode 100644 index ceb4d7df..00000000 --- a/mlx/backend/cuda/cuda.cpp +++ /dev/null @@ -1,11 +0,0 @@ -// Copyright © 2025 Apple Inc. - -#include "mlx/backend/cuda/cuda.h" - -namespace mlx::core::cu { - -bool is_available() { - return true; -} - -} // namespace mlx::core::cu diff --git a/mlx/backend/cuda/cuda.h b/mlx/backend/cuda/cuda.h index f5495965..410bae4c 100644 --- a/mlx/backend/cuda/cuda.h +++ b/mlx/backend/cuda/cuda.h @@ -2,6 +2,10 @@ #pragma once +#include +#include +#include + #include "mlx/api.h" namespace mlx::core::cu { @@ -9,4 +13,9 @@ namespace mlx::core::cu { /* Check if the CUDA backend is available. */ MLX_API bool is_available(); +/* Get information about a CUDA device. */ +MLX_API const + std::unordered_map>& + device_info(int device_index = 0); + } // namespace mlx::core::cu diff --git a/mlx/backend/cuda/device_info.cpp b/mlx/backend/cuda/device_info.cpp new file mode 100644 index 00000000..a4a15e28 --- /dev/null +++ b/mlx/backend/cuda/device_info.cpp @@ -0,0 +1,232 @@ +// Copyright © 2026 Apple Inc. + +#include "mlx/backend/gpu/device_info.h" +#include "mlx/backend/cuda/cuda.h" + +#include +#include + +#include +#include +#include +#include + +namespace mlx::core { + +namespace { + +// NVML dynamic loading for accurate memory reporting +// (cudaMemGetInfo only sees current process) + +typedef int nvmlReturn_t; +typedef struct nvmlDevice_st* nvmlDevice_t; +struct nvmlMemory_t { + unsigned long long total; + unsigned long long free; + unsigned long long used; +}; + +struct NVMLState { + void* handle = nullptr; + nvmlReturn_t (*nvmlInit_v2)() = nullptr; + nvmlReturn_t (*nvmlDeviceGetHandleByUUID)(const char*, nvmlDevice_t*) = + nullptr; + nvmlReturn_t (*nvmlDeviceGetMemoryInfo)(nvmlDevice_t, nvmlMemory_t*) = + nullptr; +}; + +bool nvml_init(NVMLState& nvml) { +#ifdef _WIN32 + nvml.handle = dlopen("nvml.dll", RTLD_LAZY); + if (!nvml.handle) { + nvml.handle = dlopen( + "C:\\Program Files\\NVIDIA Corporation\\NVSMI\\nvml.dll", RTLD_LAZY); + } +#else + nvml.handle = dlopen("libnvidia-ml.so.1", RTLD_LAZY); +#endif + if (!nvml.handle) + return false; + + nvml.nvmlInit_v2 = + (decltype(nvml.nvmlInit_v2))dlsym(nvml.handle, "nvmlInit_v2"); + nvml.nvmlDeviceGetHandleByUUID = + (decltype(nvml.nvmlDeviceGetHandleByUUID))dlsym( + nvml.handle, "nvmlDeviceGetHandleByUUID"); + nvml.nvmlDeviceGetMemoryInfo = (decltype(nvml.nvmlDeviceGetMemoryInfo))dlsym( + nvml.handle, "nvmlDeviceGetMemoryInfo"); + + if (!nvml.nvmlInit_v2 || !nvml.nvmlDeviceGetHandleByUUID || + !nvml.nvmlDeviceGetMemoryInfo) { + return false; + } + return nvml.nvmlInit_v2() == 0; +} + +bool nvml_get_memory( + NVMLState& nvml, + const char* uuid, + size_t* free, + size_t* total) { + if (!nvml.handle) + return false; + nvmlDevice_t device; + if (nvml.nvmlDeviceGetHandleByUUID(uuid, &device) != 0) + return false; + nvmlMemory_t mem; + if (nvml.nvmlDeviceGetMemoryInfo(device, &mem) != 0) + return false; + *free = mem.free; + *total = mem.total; + return true; +} + +std::string format_uuid(const cudaUUID_t& uuid) { + char buf[64]; + snprintf( + buf, + sizeof(buf), + "GPU-%02x%02x%02x%02x-%02x%02x-%02x%02x-%02x%02x-%02x%02x%02x%02x%02x%02x", + (unsigned char)uuid.bytes[0], + (unsigned char)uuid.bytes[1], + (unsigned char)uuid.bytes[2], + (unsigned char)uuid.bytes[3], + (unsigned char)uuid.bytes[4], + (unsigned char)uuid.bytes[5], + (unsigned char)uuid.bytes[6], + (unsigned char)uuid.bytes[7], + (unsigned char)uuid.bytes[8], + (unsigned char)uuid.bytes[9], + (unsigned char)uuid.bytes[10], + (unsigned char)uuid.bytes[11], + (unsigned char)uuid.bytes[12], + (unsigned char)uuid.bytes[13], + (unsigned char)uuid.bytes[14], + (unsigned char)uuid.bytes[15]); + return buf; +} + +const std::unordered_map>& +device_info_impl(int device_index) { + // Static cache of device properties including UUID (needed for NVML lookup) + static auto all_devices = []() { + // Get device count + int count = 0; + cudaGetDeviceCount(&count); + + // Collect info for all devices + struct DeviceInfo { + std::unordered_map> info; + std::string uuid; + }; + + std::vector devices; + + for (int i = 0; i < count; ++i) { + cudaDeviceProp prop; + cudaGetDeviceProperties(&prop, i); + + DeviceInfo dev; + dev.info["device_name"] = std::string(prop.name); + dev.uuid = format_uuid(prop.uuid); + dev.info["uuid"] = dev.uuid; + + // Architecture string (e.g., "sm_89") + char arch[16]; + snprintf(arch, sizeof(arch), "sm_%d%d", prop.major, prop.minor); + dev.info["architecture"] = std::string(arch); + + // PCI bus ID (domain:bus:device.function) + char pci_id[32]; + snprintf( + pci_id, + sizeof(pci_id), + "%04x:%02x:%02x.0", + prop.pciDomainID, + prop.pciBusID, + prop.pciDeviceID); + dev.info["pci_bus_id"] = std::string(pci_id); + + // Compute capability as size_t (to match Metal's variant type) + dev.info["compute_capability_major"] = static_cast(prop.major); + dev.info["compute_capability_minor"] = static_cast(prop.minor); + + devices.push_back(std::move(dev)); + } + return devices; + }(); + + // Initialize NVML once for fresh memory reads + static NVMLState nvml; + static bool nvml_initialized = nvml_init(nvml); + + if (device_index < 0 || + device_index >= static_cast(all_devices.size())) { + static auto empty = + std::unordered_map>(); + return empty; + } + + // Return a copy with fresh memory info + // Using thread_local to avoid locks while keeping free_memory fresh + thread_local auto device_info_copy = + std::unordered_map>(); + + device_info_copy = all_devices[device_index].info; + + // Get fresh memory info - try NVML first (system-wide), fallback to + // cudaMemGetInfo (process-level) + size_t free_mem, total_mem; + + if (nvml_initialized && + nvml_get_memory( + nvml, + all_devices[device_index].uuid.c_str(), + &free_mem, + &total_mem)) { + // NVML succeeded - use system-wide memory + } else { + // Fallback to cudaMemGetInfo (process-scoped) + int prev_device; + cudaGetDevice(&prev_device); + cudaSetDevice(device_index); + cudaMemGetInfo(&free_mem, &total_mem); + cudaSetDevice(prev_device); + } + + device_info_copy["free_memory"] = free_mem; + device_info_copy["total_memory"] = total_mem; + + return device_info_copy; +} + +} // anonymous namespace + +namespace gpu { + +bool is_available() { + return true; +} + +int device_count() { + int count = 0; + cudaGetDeviceCount(&count); + return count; +} + +const std::unordered_map>& +device_info(int device_index) { + return device_info_impl(device_index); +} + +} // namespace gpu + +namespace cu { + +bool is_available() { + return true; +} + +} // namespace cu + +} // namespace mlx::core diff --git a/mlx/backend/cuda/eval.cpp b/mlx/backend/cuda/eval.cpp index ef58f4a7..1d04c8f3 100644 --- a/mlx/backend/cuda/eval.cpp +++ b/mlx/backend/cuda/eval.cpp @@ -3,7 +3,6 @@ #include "mlx/backend/gpu/eval.h" #include "mlx/backend/cuda/allocator.h" #include "mlx/backend/cuda/device.h" -#include "mlx/backend/gpu/available.h" #include "mlx/primitives.h" #include "mlx/scheduler.h" @@ -11,10 +10,6 @@ namespace mlx::core::gpu { -bool is_available() { - return true; -} - void new_stream(Stream s) { // Force initalization of CUDA, so CUDA runtime get destroyed at last. cudaFree(nullptr); diff --git a/mlx/backend/gpu/available.h b/mlx/backend/gpu/available.h deleted file mode 100644 index 5d0ef7f8..00000000 --- a/mlx/backend/gpu/available.h +++ /dev/null @@ -1,11 +0,0 @@ -// Copyright © 2025 Apple Inc. - -#pragma once - -#include "mlx/api.h" - -namespace mlx::core::gpu { - -MLX_API bool is_available(); - -} // namespace mlx::core::gpu diff --git a/mlx/backend/gpu/device_info.h b/mlx/backend/gpu/device_info.h new file mode 100644 index 00000000..7adb7f0b --- /dev/null +++ b/mlx/backend/gpu/device_info.h @@ -0,0 +1,36 @@ +// Copyright © 2026 Apple Inc. + +#pragma once + +#include +#include +#include + +#include "mlx/api.h" + +namespace mlx::core::gpu { + +MLX_API bool is_available(); + +/** + * Get the number of available GPU devices. + */ +MLX_API int device_count(); + +/** + * Get information about a GPU device. + * + * Returns a map of device properties. Keys vary by backend: + * - device_name (string): Device name + * - architecture (string): Architecture identifier + * - total_memory/memory_size (size_t): Total device memory + * - free_memory (size_t): Available memory (CUDA only) + * - uuid (string): Device UUID (CUDA only) + * - pci_bus_id (string): PCI bus ID (CUDA only) + * - compute_capability_major/minor (size_t): Compute capability (CUDA only) + */ +MLX_API const + std::unordered_map>& + device_info(int device_index = 0); + +} // namespace mlx::core::gpu diff --git a/mlx/backend/metal/CMakeLists.txt b/mlx/backend/metal/CMakeLists.txt index c15f963b..c9ec78ad 100644 --- a/mlx/backend/metal/CMakeLists.txt +++ b/mlx/backend/metal/CMakeLists.txt @@ -102,6 +102,7 @@ target_sources( mlx PRIVATE ${CMAKE_CURRENT_SOURCE_DIR}/allocator.cpp ${CMAKE_CURRENT_SOURCE_DIR}/binary.cpp + ${CMAKE_CURRENT_SOURCE_DIR}/device_info.cpp ${CMAKE_CURRENT_SOURCE_DIR}/compiled.cpp ${CMAKE_CURRENT_SOURCE_DIR}/conv.cpp ${CMAKE_CURRENT_SOURCE_DIR}/copy.cpp diff --git a/mlx/backend/metal/allocator.cpp b/mlx/backend/metal/allocator.cpp index 2e7fea2a..06e3d9ae 100644 --- a/mlx/backend/metal/allocator.cpp +++ b/mlx/backend/metal/allocator.cpp @@ -1,5 +1,6 @@ // Copyright © 2023-2024 Apple Inc. #include "mlx/backend/metal/allocator.h" +#include "mlx/backend/gpu/device_info.h" #include "mlx/backend/metal/metal.h" #include "mlx/backend/metal/resident.h" #include "mlx/memory.h" @@ -43,16 +44,17 @@ MetalAllocator::MetalAllocator() }), residency_set_(device_) { auto pool = metal::new_scoped_memory_pool(); - auto memsize = std::get(device_info().at("memory_size")); + const auto& info = gpu::device_info(0); + auto memsize = std::get(info.at("memory_size")); auto max_rec_size = - std::get(device_info().at("max_recommended_working_set_size")); - resource_limit_ = std::get(device_info().at("resource_limit")); + std::get(info.at("max_recommended_working_set_size")); + resource_limit_ = std::get(info.at("resource_limit")); block_limit_ = std::min(1.5 * max_rec_size, 0.95 * memsize); gc_limit_ = std::min(static_cast(0.95 * max_rec_size), block_limit_); max_pool_size_ = block_limit_; device(mlx::core::Device::gpu) .set_residency_set(residency_set_.mtl_residency_set()); - bool is_vm = std::get(device_info().at("device_name")) == + bool is_vm = std::get(info.at("device_name")) == "Apple Paravirtual device"; if (is_vm) { return; @@ -249,8 +251,8 @@ size_t get_memory_limit() { return metal::allocator().get_memory_limit(); } size_t set_wired_limit(size_t limit) { - if (limit > std::get(metal::device_info().at( - "max_recommended_working_set_size"))) { + if (limit > std::get( + gpu::device_info(0).at("max_recommended_working_set_size"))) { throw std::invalid_argument( "[metal::set_wired_limit] Setting a wired limit larger than " "the maximum working set size is not allowed."); diff --git a/mlx/backend/metal/device_info.cpp b/mlx/backend/metal/device_info.cpp new file mode 100644 index 00000000..dd18fc6c --- /dev/null +++ b/mlx/backend/metal/device_info.cpp @@ -0,0 +1,58 @@ +// Copyright © 2026 Apple Inc. + +#include + +#include "mlx/backend/gpu/device_info.h" +#include "mlx/backend/metal/device.h" +#include "mlx/backend/metal/metal.h" + +namespace mlx::core::gpu { + +bool is_available() { + return metal::is_available(); +} + +int device_count() { + return 1; +} + +const std::unordered_map>& +device_info(int device_index) { + auto init_device_info = []() + -> std::unordered_map> { + auto pool = metal::new_scoped_memory_pool(); + auto raw_device = metal::device(mlx::core::Device::gpu).mtl_device(); + auto name = std::string(raw_device->name()->utf8String()); + auto arch = std::string(raw_device->architecture()->name()->utf8String()); + + size_t memsize = 0; + size_t length = sizeof(memsize); + sysctlbyname("hw.memsize", &memsize, &length, NULL, 0); + + size_t rsrc_limit = 0; + sysctlbyname("iogpu.rsrc_limit", &rsrc_limit, &length, NULL, 0); + if (rsrc_limit == 0) { + rsrc_limit = 499000; + } + + return { + {"device_name", name}, + {"architecture", arch}, + {"max_buffer_length", raw_device->maxBufferLength()}, + {"max_recommended_working_set_size", + raw_device->recommendedMaxWorkingSetSize()}, + {"memory_size", memsize}, + {"resource_limit", rsrc_limit}}; + }; + static auto device_info_ = init_device_info(); + static std::unordered_map> + empty; + + if (device_index == 0) { + return device_info_; + } else { + return empty; + } +} + +} // namespace mlx::core::gpu diff --git a/mlx/backend/metal/eval.cpp b/mlx/backend/metal/eval.cpp index 10005892..bd58a691 100644 --- a/mlx/backend/metal/eval.cpp +++ b/mlx/backend/metal/eval.cpp @@ -1,7 +1,6 @@ // Copyright © 2023-2024 Apple Inc. #include -#include "mlx/backend/gpu/available.h" #include "mlx/backend/gpu/eval.h" #include "mlx/backend/metal/device.h" #include "mlx/backend/metal/utils.h" @@ -10,10 +9,6 @@ namespace mlx::core::gpu { -bool is_available() { - return true; -} - void new_stream(Stream stream) { if (stream.device == mlx::core::Device::gpu) { metal::device(stream.device).new_queue(stream.index); diff --git a/mlx/backend/metal/metal.cpp b/mlx/backend/metal/metal.cpp index 078ea70d..51bd2e62 100644 --- a/mlx/backend/metal/metal.cpp +++ b/mlx/backend/metal/metal.cpp @@ -1,8 +1,6 @@ // Copyright © 2023-2024 Apple Inc. #include -#include - #include "mlx/backend/metal/device.h" #include "mlx/backend/metal/metal.h" #include "mlx/backend/metal/utils.h" @@ -49,36 +47,4 @@ void stop_capture() { manager->stopCapture(); } -const std::unordered_map>& -device_info() { - auto init_device_info = []() - -> std::unordered_map> { - auto pool = new_scoped_memory_pool(); - auto raw_device = device(default_device()).mtl_device(); - auto name = std::string(raw_device->name()->utf8String()); - auto arch = std::string(raw_device->architecture()->name()->utf8String()); - - size_t memsize = 0; - size_t length = sizeof(memsize); - sysctlbyname("hw.memsize", &memsize, &length, NULL, 0); - - size_t rsrc_limit = 0; - sysctlbyname("iogpu.rsrc_limit", &rsrc_limit, &length, NULL, 0); - if (rsrc_limit == 0) { - rsrc_limit = 499000; - } - - return { - {"device_name", name}, - {"architecture", arch}, - {"max_buffer_length", raw_device->maxBufferLength()}, - {"max_recommended_working_set_size", - raw_device->recommendedMaxWorkingSetSize()}, - {"memory_size", memsize}, - {"resource_limit", rsrc_limit}}; - }; - static auto device_info_ = init_device_info(); - return device_info_; -} - } // namespace mlx::core::metal diff --git a/mlx/backend/no_cpu/CMakeLists.txt b/mlx/backend/no_cpu/CMakeLists.txt index 2e696082..c81a2585 100644 --- a/mlx/backend/no_cpu/CMakeLists.txt +++ b/mlx/backend/no_cpu/CMakeLists.txt @@ -1,6 +1,6 @@ target_sources( mlx - PRIVATE ${CMAKE_CURRENT_SOURCE_DIR}/available.cpp + PRIVATE ${CMAKE_CURRENT_SOURCE_DIR}/device_info.cpp ${CMAKE_CURRENT_SOURCE_DIR}/primitives.cpp ${CMAKE_CURRENT_SOURCE_DIR}/../cpu/eval.cpp ${CMAKE_CURRENT_SOURCE_DIR}/../cpu/encoder.cpp diff --git a/mlx/backend/no_cpu/available.cpp b/mlx/backend/no_cpu/available.cpp deleted file mode 100644 index 04c1bac8..00000000 --- a/mlx/backend/no_cpu/available.cpp +++ /dev/null @@ -1,11 +0,0 @@ -// Copyright © 2025 Apple Inc. - -#include "mlx/backend/cpu/available.h" - -namespace mlx::core::cpu { - -bool is_available() { - return false; -} - -} // namespace mlx::core::cpu diff --git a/mlx/backend/no_cpu/device_info.cpp b/mlx/backend/no_cpu/device_info.cpp new file mode 100644 index 00000000..0fc07fc3 --- /dev/null +++ b/mlx/backend/no_cpu/device_info.cpp @@ -0,0 +1,22 @@ +// Copyright © 2026 Apple Inc. + +#include "mlx/backend/cpu/device_info.h" + +namespace mlx::core::cpu { + +bool is_available() { + return false; +} + +int device_count() { + return 0; +} + +const std::unordered_map>& +device_info(int /* device_index */) { + static std::unordered_map> + empty; + return empty; +} + +} // namespace mlx::core::cpu diff --git a/mlx/backend/no_gpu/CMakeLists.txt b/mlx/backend/no_gpu/CMakeLists.txt index 78e15ac6..c777204d 100644 --- a/mlx/backend/no_gpu/CMakeLists.txt +++ b/mlx/backend/no_gpu/CMakeLists.txt @@ -1,6 +1,7 @@ target_sources( mlx PRIVATE ${CMAKE_CURRENT_SOURCE_DIR}/allocator.cpp + ${CMAKE_CURRENT_SOURCE_DIR}/device_info.cpp ${CMAKE_CURRENT_SOURCE_DIR}/event.cpp ${CMAKE_CURRENT_SOURCE_DIR}/fence.cpp ${CMAKE_CURRENT_SOURCE_DIR}/eval.cpp diff --git a/mlx/backend/no_gpu/device_info.cpp b/mlx/backend/no_gpu/device_info.cpp new file mode 100644 index 00000000..65d1cfcc --- /dev/null +++ b/mlx/backend/no_gpu/device_info.cpp @@ -0,0 +1,22 @@ +// Copyright © 2026 Apple Inc. + +#include "mlx/backend/gpu/device_info.h" + +namespace mlx::core::gpu { + +bool is_available() { + return false; +} + +int device_count() { + return 0; +} + +const std::unordered_map>& +device_info(int /* device_index */) { + static std::unordered_map> + empty; + return empty; +} + +} // namespace mlx::core::gpu diff --git a/mlx/backend/no_gpu/eval.cpp b/mlx/backend/no_gpu/eval.cpp index 8bff86a9..5d88970a 100644 --- a/mlx/backend/no_gpu/eval.cpp +++ b/mlx/backend/no_gpu/eval.cpp @@ -2,15 +2,11 @@ #include -#include "mlx/backend/gpu/available.h" +#include "mlx/backend/gpu/device_info.h" #include "mlx/backend/gpu/eval.h" namespace mlx::core::gpu { -bool is_available() { - return false; -} - void new_stream(Stream) {} void eval(array&) { diff --git a/mlx/device.cpp b/mlx/device.cpp index ec17a509..4a62036e 100644 --- a/mlx/device.cpp +++ b/mlx/device.cpp @@ -1,9 +1,9 @@ -// Copyright © 2023 Apple Inc. +// Copyright © 2023-2026 Apple Inc. #include -#include "mlx/backend/cpu/available.h" -#include "mlx/backend/gpu/available.h" +#include "mlx/backend/cpu/device_info.h" +#include "mlx/backend/gpu/device_info.h" #include "mlx/device.h" namespace mlx::core { @@ -44,4 +44,29 @@ bool is_available(const Device& d) { return false; } +int device_count(Device::DeviceType type) { + switch (type) { + case Device::cpu: + return cpu::device_count(); + case Device::gpu: + return gpu::device_count(); + } + // appease compiler + return 0; +} + +const std::unordered_map>& +device_info(const Device& d) { + switch (d.type) { + case Device::cpu: + return cpu::device_info(d.index); + case Device::gpu: + return gpu::device_info(d.index); + } + // appease compiler + static std::unordered_map> + empty; + return empty; +} + } // namespace mlx::core diff --git a/mlx/device.h b/mlx/device.h index fb14b917..f89ad189 100644 --- a/mlx/device.h +++ b/mlx/device.h @@ -1,9 +1,13 @@ -// Copyright © 2023 Apple Inc. +// Copyright © 2023-2025 Apple Inc. #pragma once #include "mlx/api.h" +#include +#include +#include + namespace mlx::core { struct MLX_API Device { @@ -30,4 +34,23 @@ MLX_API bool operator!=(const Device& lhs, const Device& rhs); MLX_API bool is_available(const Device& d); +/** Get the number of available devices for the given device type. */ +MLX_API int device_count(Device::DeviceType type); + +/** + * Get information about a device. + * + * Returns a map of device properties. Keys vary by backend: + * - device_name (string): Device name + * - architecture (string): Architecture identifier + * - total_memory/memory_size (size_t): Total device memory + * - free_memory (size_t): Available memory (CUDA only) + * - uuid (string): Device UUID (CUDA only) + * - pci_bus_id (string): PCI bus ID (CUDA only) + * - compute_capability_major/minor (size_t): Compute capability (CUDA only) + */ +MLX_API const + std::unordered_map>& + device_info(const Device& d = default_device()); + } // namespace mlx::core diff --git a/mlx/mlx.h b/mlx/mlx.h index dbc9014d..eda7333d 100644 --- a/mlx/mlx.h +++ b/mlx/mlx.h @@ -4,7 +4,7 @@ #include "mlx/array.h" #include "mlx/backend/cuda/cuda.h" -#include "mlx/backend/gpu/available.h" +#include "mlx/backend/gpu/device_info.h" #include "mlx/backend/metal/metal.h" #include "mlx/compile.h" #include "mlx/device.h" diff --git a/mlx/scheduler.cpp b/mlx/scheduler.cpp index b19f6434..87ff98e2 100644 --- a/mlx/scheduler.cpp +++ b/mlx/scheduler.cpp @@ -1,7 +1,7 @@ // Copyright © 2023 Apple Inc. #include "mlx/scheduler.h" -#include "mlx/backend/gpu/available.h" +#include "mlx/backend/gpu/device_info.h" #include "mlx/backend/gpu/eval.h" namespace mlx::core { diff --git a/python/src/device.cpp b/python/src/device.cpp index 006a05dc..f15f7f92 100644 --- a/python/src/device.cpp +++ b/python/src/device.cpp @@ -1,9 +1,11 @@ -// Copyright © 2023-2024 Apple Inc. +// Copyright © 2023-2025 Apple Inc. #include #include #include +#include +#include #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"); } diff --git a/python/src/metal.cpp b/python/src/metal.cpp index a5667442..fb3c5130 100644 --- a/python/src/metal.cpp +++ b/python/src/metal.cpp @@ -9,6 +9,7 @@ #include #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)); + }); } diff --git a/python/tests/test_device.py b/python/tests/test_device.py index d51028de..6cd6df81 100644 --- a/python/tests/test_device.py +++ b/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() diff --git a/python/tests/test_memory.py b/python/tests/test_memory.py index da4a238d..2a0a544b 100644 --- a/python/tests/test_memory.py +++ b/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)