diff --git a/docs/src/python/cuda.rst b/docs/src/python/cuda.rst index 932d36b5ef..4ffefa1da3 100644 --- a/docs/src/python/cuda.rst +++ b/docs/src/python/cuda.rst @@ -7,3 +7,4 @@ CUDA :toctree: _autosummary is_available + get_arch_list diff --git a/mlx/backend/cuda/CMakeLists.txt b/mlx/backend/cuda/CMakeLists.txt index 50be6f404f..0013892714 100644 --- a/mlx/backend/cuda/CMakeLists.txt +++ b/mlx/backend/cuda/CMakeLists.txt @@ -181,10 +181,38 @@ if(NOT DEFINED MLX_CUDA_ARCHITECTURES) set(MLX_CUDA_ARCHITECTURES "${MLX_CUDA_ARCHITECTURES}a") endif() endif() +# Resolve all/all-major/native aliases. +if(MLX_CUDA_ARCHITECTURES STREQUAL "all") + set(MLX_CUDA_ARCHITECTURES "${CMAKE_CUDA_ARCHITECTURES_ALL}") +elseif(MLX_CUDA_ARCHITECTURES STREQUAL "all-major") + set(MLX_CUDA_ARCHITECTURES "${CMAKE_CUDA_ARCHITECTURES_ALL_MAJOR}") +elseif(MLX_CUDA_ARCHITECTURES STREQUAL "native") + set(MLX_CUDA_ARCHITECTURES "${CMAKE_CUDA_ARCHITECTURES_NATIVE}") +endif() message(STATUS "CUDA architectures: ${MLX_CUDA_ARCHITECTURES}") set_target_properties(mlx PROPERTIES CUDA_ARCHITECTURES "${MLX_CUDA_ARCHITECTURES}") +# Pass the list of cuda arch to code. +set(MLX_CUDA_ARCH_LIST) +foreach(arch IN LISTS MLX_CUDA_ARCHITECTURES) + if(arch MATCHES "^([0-9]+[af]?)-(real|virtual)$") + if(CMAKE_MATCH_2 STREQUAL "real") + list(APPEND MLX_CUDA_ARCH_LIST "\"sm_${CMAKE_MATCH_1}\"") + else() + list(APPEND MLX_CUDA_ARCH_LIST "\"compute_${CMAKE_MATCH_1}\"") + endif() + elseif(arch MATCHES "^[0-9]+[af]?$") + list(APPEND MLX_CUDA_ARCH_LIST "\"sm_${arch}\"" "\"compute_${arch}\"") + else() + list(APPEND MLX_CUDA_ARCH_LIST "\"${arch}\"") + endif() +endforeach() +list(JOIN MLX_CUDA_ARCH_LIST "," MLX_CUDA_ARCH_LIST) +target_compile_definitions( + mlx + PRIVATE "$<$:MLX_CUDA_ARCH_LIST=${MLX_CUDA_ARCH_LIST}>") + # Configure Windows CUDA DLL loading. if(WIN32) target_sources(mlx PRIVATE ${CMAKE_CURRENT_SOURCE_DIR}/wddm.cpp) diff --git a/mlx/backend/cuda/cuda.h b/mlx/backend/cuda/cuda.h index 410bae4c63..d6b8b73fef 100644 --- a/mlx/backend/cuda/cuda.h +++ b/mlx/backend/cuda/cuda.h @@ -5,6 +5,7 @@ #include #include #include +#include #include "mlx/api.h" @@ -13,6 +14,9 @@ namespace mlx::core::cu { /* Check if the CUDA backend is available. */ MLX_API bool is_available(); +/* Get the CUDA architectures MLX was compiled for. */ +MLX_API std::vector get_arch_list(); + /* Get information about a CUDA device. */ MLX_API const std::unordered_map>& diff --git a/mlx/backend/cuda/device_info.cpp b/mlx/backend/cuda/device_info.cpp index a4a15e2850..43c25ee8cb 100644 --- a/mlx/backend/cuda/device_info.cpp +++ b/mlx/backend/cuda/device_info.cpp @@ -227,6 +227,10 @@ bool is_available() { return true; } +std::vector get_arch_list() { + return {MLX_CUDA_ARCH_LIST}; +} + } // namespace cu } // namespace mlx::core diff --git a/mlx/backend/cuda/no_cuda.cpp b/mlx/backend/cuda/no_cuda.cpp index 175a505b4b..8126790dd4 100644 --- a/mlx/backend/cuda/no_cuda.cpp +++ b/mlx/backend/cuda/no_cuda.cpp @@ -11,6 +11,10 @@ bool is_available() { return false; } +std::vector get_arch_list() { + return {}; +} + } // namespace cu namespace fast { diff --git a/python/src/cuda.cpp b/python/src/cuda.cpp index 13b3a0154b..34dec0fc69 100644 --- a/python/src/cuda.cpp +++ b/python/src/cuda.cpp @@ -1,6 +1,8 @@ // Copyright © 2023-2025 Apple Inc. #include +#include +#include #include "mlx/backend/cuda/cuda.h" @@ -16,4 +18,17 @@ void init_cuda(nb::module_& m) { R"pbdoc( Check if the CUDA back-end is available. )pbdoc"); + cuda.def( + "get_arch_list", + &mx::cu::get_arch_list, + R"pbdoc( + Get the CUDA architectures MLX was compiled for. + + Architectures use NVIDIA's ``sm_*`` (binary) and ``compute_*`` (PTX) + notation, for example ``["sm_80", "sm_120a", "compute_120"]``. + + Returns: + list(str): The compiled architectures, or an empty list if CUDA + backend is not enabled. + )pbdoc"); } diff --git a/python/tests/test_device.py b/python/tests/test_device.py index bef788bb1d..86f7fbebd1 100644 --- a/python/tests/test_device.py +++ b/python/tests/test_device.py @@ -184,6 +184,16 @@ def test_device_info_gpu(self): self.assertTrue(len(info["device_name"]) > 0) self.assertIn("architecture", info) + def test_cuda_get_arch_list(self): + arch_list = mx.cuda.get_arch_list() + self.assertIsInstance(arch_list, list) + if mx.cuda.is_available(): + self.assertTrue(arch_list) + for arch in arch_list: + self.assertRegex(arch, r"^(sm|compute)_[0-9]+[af]?$") + else: + self.assertEqual(arch_list, []) + def test_device_info_default(self): info = mx.device_info() self.assertIsInstance(info, dict)