fix(moore): select MCCL bfloat16 architecture

CMakeLists.txt

@@ -1,6 +1,8 @@
 cmake_minimum_required(VERSION 3.18)
 project(InfiniCCL VERSION 0.1.0 LANGUAGES C CXX)
 
+include(cmake/MooreArchitecture.cmake)
+
 set(CMAKE_CXX_STANDARD 17)
 set(CMAKE_CXX_STANDARD_REQUIRED ON)
 
@@ -31,6 +33,9 @@ option(AUTO_DETECT_DEVICES "Automatically detect available devices" ON)
 option(AUTO_DETECT_BACKENDS "Automatically detect available backends" OFF)
 option(BUILD_EXAMPLES "Build internal examples" ON)
 
+set(MUSA_ARCHITECTURES "" CACHE STRING
+    "MUSA GPU architectures (for example `31;22`); detected from installed GPUs when empty")
+
 # =========================================================
 # --- AUTO-DETECTION: DEVICES ---
 # =========================================================
@@ -397,6 +402,34 @@ if(WITH_MOORE)
     find_library(MUSA_LIB   NAMES musa   HINTS "${MUSA_ROOT}/lib" REQUIRED)
     find_library(MUSART_LIB NAMES musart HINTS "${MUSA_ROOT}/lib" REQUIRED)
     find_library(MUBLAS_LIB NAMES mublas HINTS "${MUSA_ROOT}/lib" REQUIRED)
+
+    set(_musa_architectures "${MUSA_ARCHITECTURES}")
+    if(NOT _musa_architectures
+       AND DEFINED ENV{TORCH_MUSA_ARCH_LIST}
+       AND NOT "$ENV{TORCH_MUSA_ARCH_LIST}" STREQUAL "")
+        set(_musa_architectures "$ENV{TORCH_MUSA_ARCH_LIST}")
+        message(STATUS "Using MUSA architectures from `TORCH_MUSA_ARCH_LIST`.")
+    endif()
+
+    if(NOT _musa_architectures)
+        infiniccl_detect_musa_architectures(
+            "${MUSA_ROOT}/include"
+            "${MUSART_LIB}"
+            _musa_architectures
+        )
+        message(STATUS "Auto-detected MUSA architectures from installed GPUs.")
+    endif()
+
+    infiniccl_compute_musa_architecture_config(
+        "${_musa_architectures}"
+        MUSA_ARCHITECTURES
+        MUSA_MARCH_TYPE
+        MUSA_ARCH_COMPILE_OPTIONS
+    )
+    set(MUSA_ARCHITECTURES "${MUSA_ARCHITECTURES}" CACHE STRING
+        "MUSA GPU architectures (for example `31;22`); detected from installed GPUs when empty" FORCE)
+    message(STATUS
+        "MUSA architectures: ${MUSA_ARCHITECTURES} (`MARCH_TYPE=${MUSA_MARCH_TYPE}`)")
 endif()
 
 if(WITH_CAMBRICON)
@@ -507,7 +540,6 @@ add_subdirectory(src)
 if(BUILD_EXAMPLES)
     add_subdirectory(examples)
 endif()
-
 # =========================================================
 # --- Installation ---
 # =========================================================

cmake/MooreArchitecture.cmake

@@ -0,0 +1,115 @@
+function(infiniccl_normalize_musa_architectures input output)
+    set(_architectures "${input}")
+    string(REPLACE "," ";" _architectures "${_architectures}")
+    string(REGEX REPLACE "[ \t\r\n]+" ";" _architectures "${_architectures}")
+
+    set(_normalized)
+    foreach(_architecture IN LISTS _architectures)
+        if(_architecture STREQUAL "")
+            continue()
+        endif()
+
+        string(REGEX REPLACE "^mp_" "" _architecture "${_architecture}")
+        string(REPLACE "." "" _architecture "${_architecture}")
+        if(NOT _architecture MATCHES "^[0-9]+$")
+            message(FATAL_ERROR
+                "Invalid MUSA architecture `${_architecture}`. "
+                "Use values such as `31`, `mp_31`, or `3.1`."
+            )
+        endif()
+
+        list(APPEND _normalized "${_architecture}")
+    endforeach()
+
+    list(REMOVE_DUPLICATES _normalized)
+    if(NOT _normalized)
+        message(FATAL_ERROR "At least one MUSA architecture is required.")
+    endif()
+
+    set(${output} "${_normalized}" PARENT_SCOPE)
+endfunction()
+
+function(infiniccl_compute_musa_architecture_config input architectures_output march_type_output flags_output)
+    infiniccl_normalize_musa_architectures("${input}" _architectures)
+
+    # `mccl.h` exposes data types globally from one `MARCH_TYPE`. Use the least
+    # capable target so a fat binary never advertises a type that one of its
+    # target architectures cannot execute.
+    set(_march_type)
+    set(_flags)
+    foreach(_architecture IN LISTS _architectures)
+        math(EXPR _architecture_version "${_architecture} * 10")
+        if(NOT _march_type OR _architecture_version LESS _march_type)
+            set(_march_type "${_architecture_version}")
+        endif()
+        list(APPEND _flags "--offload-arch=mp_${_architecture}")
+    endforeach()
+
+    set(${architectures_output} "${_architectures}" PARENT_SCOPE)
+    set(${march_type_output} "${_march_type}" PARENT_SCOPE)
+    set(${flags_output} "${_flags}" PARENT_SCOPE)
+endfunction()
+
+function(infiniccl_configure_musa_target target)
+    if(NOT TARGET ${target})
+        message(FATAL_ERROR "Cannot configure unknown MUSA target `${target}`.")
+    endif()
+
+    foreach(_musa_arch_compile_option IN LISTS MUSA_ARCH_COMPILE_OPTIONS)
+        target_compile_options(${target} PRIVATE
+            $<$<COMPILE_LANGUAGE:CXX>:${_musa_arch_compile_option}>
+        )
+    endforeach()
+    target_compile_definitions(${target} PRIVATE
+        MARCH_TYPE=${MUSA_MARCH_TYPE}
+    )
+endfunction()
+
+function(infiniccl_detect_musa_architectures musa_include_dir musart_library output)
+    set(_source "${CMAKE_CURRENT_BINARY_DIR}/get_musa_compute_capabilities.cpp")
+    file(WRITE "${_source}" [=[
+#include <musa_runtime.h>
+
+#include <cstdio>
+
+int main() {
+  int device_count = 0;
+  if (musaGetDeviceCount(&device_count) != musaSuccess || device_count == 0) {
+    return 1;
+  }
+
+  for (int device = 0; device < device_count; ++device) {
+    musaDeviceProp properties;
+    if (musaGetDeviceProperties(&properties, device) != musaSuccess) {
+      return 1;
+    }
+    std::printf("%d%d ", properties.major, properties.minor);
+  }
+  return 0;
+}
+]=])
+
+    try_run(
+        _run_result
+        _compile_result
+        "${CMAKE_CURRENT_BINARY_DIR}"
+        "${_source}"
+        CMAKE_FLAGS "-DINCLUDE_DIRECTORIES=${musa_include_dir}"
+        LINK_LIBRARIES "${musart_library}"
+        RUN_OUTPUT_VARIABLE _detected_architectures
+    )
+
+    if(NOT _compile_result OR NOT "${_run_result}" STREQUAL "0")
+        message(FATAL_ERROR
+            "Could not detect the architecture of the installed MUSA GPUs. "
+            "Set `MUSA_ARCHITECTURES` explicitly, for example "
+            "`-DMUSA_ARCHITECTURES=31`."
+        )
+    endif()
+
+    infiniccl_normalize_musa_architectures(
+        "${_detected_architectures}"
+        _detected_architectures
+    )
+    set(${output} "${_detected_architectures}" PARENT_SCOPE)
+endfunction()

examples/CMakeLists.txt

@@ -43,6 +43,7 @@ foreach(source_file ${EXAMPLE_SOURCES})
     if(WITH_MOORE)
         target_link_libraries(${target_name} PRIVATE ${MUSART_LIB})
         target_compile_options(${target_name} PRIVATE "-x" "musa")
+        infiniccl_configure_musa_target(${target_name})
     endif()
 
     if(WITH_CAMBRICON)

src/CMakeLists.txt

@@ -121,6 +121,7 @@ if(WITH_MOORE)
     target_include_directories(infiniccl PRIVATE "${MUSA_ROOT}/include")
     target_link_libraries(infiniccl PRIVATE ${MUSA_LIB} ${MUSART_LIB} ${MUBLAS_LIB})
     target_compile_options(infiniccl PRIVATE $<$<COMPILE_LANGUAGE:CXX>:-x musa>)
+    infiniccl_configure_musa_target(infiniccl)
 endif()
 
 # Cambricon