[WIP] feat: add MLU backend support

backends/mlu/CMakeLists.txt

@@ -0,0 +1,77 @@
+# TODO(mlu): CNCL distributed collectives land with the kernels milestone.
+find_library(MLU_RUNTIME_LIBRARY NAMES cnrt HINTS "${NEUWARE_HOME}/lib64" REQUIRED)
+find_library(MLU_DNN_LIBRARY NAMES cnnl HINTS "${NEUWARE_HOME}/lib64" REQUIRED)
+
+set(MLU_BACKEND_SOURCES
+  src/backend.cc
+  src/kernels/cast.cc
+  src/kernels/fill.cc
+  src/kernels/no_op.cc
+  src/runtime/mlu_guard_impl.cc
+  src/runtime/mlu_runtime_common.cc
+)
+
+# BANG device kernels, compiled with the vendor toolchain; the object carries
+# both the device code and its host launch wrappers.
+add_custom_command(OUTPUT ${CMAKE_CURRENT_BINARY_DIR}/mlu_fill_device.o
+  COMMAND ${NEUWARE_HOME}/bin/cncc --bang-arch=compute_50 -std=c++14 -fPIE -c ${CMAKE_CURRENT_SOURCE_DIR}/src/kernels/fill.mlu
+          -o ${CMAKE_CURRENT_BINARY_DIR}/mlu_fill_device.o
+  DEPENDS src/kernels/fill.mlu
+  COMMENT "cncc: compiling src/kernels/fill.mlu for compute_50"
+  VERBATIM)
+
+add_library(infini_train_backend_mlu STATIC ${MLU_BACKEND_SOURCES} ${CMAKE_CURRENT_BINARY_DIR}/mlu_fill_device.o)
+add_library(InfiniTrain::Backend::MLU ALIAS infini_train_backend_mlu)
+
+target_include_directories(infini_train_backend_mlu
+  PUBLIC
+    "$<BUILD_INTERFACE:${CMAKE_CURRENT_SOURCE_DIR}/include>"
+  PRIVATE
+    "${CMAKE_CURRENT_SOURCE_DIR}/src"
+    "${INFINITRAIN_SOURCE_DIR}"
+    "${NEUWARE_HOME}/include"
+)
+
+target_link_libraries(infini_train_backend_mlu PUBLIC
+  InfiniTrain::infini_train
+  ${MLU_RUNTIME_LIBRARY}
+  ${MLU_DNN_LIBRARY}
+)
+
+# Static registries in InfiniTrain and its CPU kernel library are discovered by
+# global initialization. Executables should link this target so those objects
+# are retained even when no ordinary symbol references them directly.
+add_library(infini_train_backend_mlu_executable INTERFACE)
+add_library(InfiniTrain::Backend::MLUExecutable ALIAS
+  infini_train_backend_mlu_executable)
+target_link_libraries(infini_train_backend_mlu_executable INTERFACE
+  "-Wl,--start-group"
+  "-Wl,--whole-archive"
+  infini_train
+  infini_train_cpu_kernels
+  infini_train_backend_mlu
+  "-Wl,--no-whole-archive"
+  "-Wl,--end-group"
+)
+
+if(BUILD_TEST)
+  if(NOT COMMAND infini_train_add_privateuse1_test_suites)
+    message(FATAL_ERROR
+      "The selected InfiniTrain revision does not provide PrivateUse1 test suites. "
+      "Update INFINITRAIN_SOURCE_DIR or the pinned submodule revision.")
+  endif()
+
+  infini_train_add_privateuse1_test_suites(
+    BACKEND_NAME ${INFINITRAIN_BACKEND}
+    DEVICE_INDEX 0
+    LINK_LIBRARIES InfiniTrain::Backend::MLUExecutable
+    BACKEND_HEADER "infini_train_mlu/backend.h"
+    BACKEND_REGISTRAR infini_train::mlu::RegisterBackend
+    TEST_TIMEOUT 300
+    RUN_SERIAL
+  )
+endif()
+
+if(INFINITRAIN_BACKENDS_BUILD_EXAMPLES)
+  add_subdirectory(examples)
+endif()

backends/mlu/cmake/pre_project.cmake

@@ -0,0 +1,20 @@
+# MLU provider pre-project setup.
+#
+# Unlike MACA (whose mxcc serves as the whole-build compiler), the InfiniTrain
+# core and this provider's host-side code compile with the system C++20
+# compiler. The vendor compiler cncc is reserved for device kernels (.mlu),
+# which will be introduced together with the kernels.
+set(NEUWARE_HOME "$ENV{NEUWARE_HOME}" CACHE PATH "Path to the Neuware SDK")
+if(NOT NEUWARE_HOME)
+  set(NEUWARE_HOME "/usr/local/neuware" CACHE PATH "Path to the Neuware SDK" FORCE)
+endif()
+if(NOT EXISTS "${NEUWARE_HOME}/include/cnrt.h")
+  message(FATAL_ERROR
+    "cnrt.h was not found under ${NEUWARE_HOME}/include. Export NEUWARE_HOME "
+    "or pass -DNEUWARE_HOME=<sdk-root>.")
+endif()
+
+set(USE_CUDA OFF)
+set(USE_NCCL OFF)
+set(USE_OMP ON)
+set(BUILD_SHARED_LIBS OFF)

backends/mlu/examples/CMakeLists.txt

@@ -0,0 +1,36 @@
+function(add_mlu_example target_name output_name)
+  add_executable(${target_name} ${ARGN})
+  # Reuse upstream example sources while injecting the selected external
+  # provider's declaration and registration callback.
+  target_compile_definitions(${target_name} PRIVATE
+    INFINITRAIN_EXAMPLE_EXTERNAL_BACKEND_HEADER=\"infini_train_mlu/backend.h\"
+    INFINITRAIN_EXAMPLE_EXTERNAL_BACKEND_REGISTRAR=infini_train::mlu::RegisterBackend
+  )
+  target_link_libraries(${target_name} PRIVATE
+    InfiniTrain::Backend::MLUExecutable)
+  set_target_properties(${target_name} PROPERTIES
+    OUTPUT_NAME ${output_name}
+    RUNTIME_OUTPUT_DIRECTORY "${CMAKE_BINARY_DIR}"
+  )
+endfunction()
+
+function(add_mlu_model_example model_name)
+  set(model_dir "${INFINITRAIN_SOURCE_DIR}/example/${model_name}")
+  set(model_sources
+    "${INFINITRAIN_SOURCE_DIR}/example/common/tiny_shakespeare_dataset.cc"
+    "${INFINITRAIN_SOURCE_DIR}/example/common/utils.cc"
+    "${model_dir}/checkpoint_loader.cc"
+  )
+  if(NOT "${model_name}" STREQUAL "mixtral")
+    list(APPEND model_sources
+      "${INFINITRAIN_SOURCE_DIR}/example/common/tokenizer.cc")
+  endif()
+
+  list(APPEND model_sources "${model_dir}/main.cc")
+
+  add_mlu_example("mlu_${model_name}" "${model_name}" ${model_sources})
+endfunction()
+
+add_mlu_model_example(gpt2)
+add_mlu_model_example(llama3)
+add_mlu_model_example(mixtral)

backends/mlu/include/infini_train_mlu/backend.h

@@ -0,0 +1,9 @@
+#pragma once
+
+namespace infini_train::mlu {
+
+// Registers MLU as the process-wide PrivateUse1 provider. Device runtime
+// initialization remains lazy until the first device operation.
+void RegisterBackend();
+
+} // namespace infini_train::mlu

backends/mlu/src/backend.cc

@@ -0,0 +1,19 @@
+#include "infini_train_mlu/backend.h"
+
+#include <mutex>
+
+#include "infini_train/include/core/privateuse1_backend.h"
+
+namespace infini_train::mlu {
+
+void RegisterBackend() {
+    static std::once_flag once;
+    std::call_once(once, []() {
+        core::PrivateUse1BackendRegistration registration;
+        registration.name = "mlu";
+        registration.default_autocast_dtype = DataType::kBFLOAT16;
+        core::RegisterPrivateUse1Backend(registration);
+    });
+}
+
+} // namespace infini_train::mlu

backends/mlu/src/common/common_mlu.h

@@ -0,0 +1,24 @@
+#pragma once
+
+#include <cnnl.h>
+#include <cnrt.h>
+
+#include "glog/logging.h"
+
+namespace infini_train::common::mlu {
+
+// Runs a CNRT call and aborts with the error string on failure.
+#define MLU_CHECK(expr)                                                                                                \
+    do {                                                                                                               \
+        cnrtRet_t _mlu_ret = (expr);                                                                                   \
+        CHECK_EQ(_mlu_ret, cnrtSuccess) << "CNRT call failed: " #expr " -> " << cnrtGetErrorStr(_mlu_ret);             \
+    } while (0)
+
+// Runs a CNNL call and aborts with the status code on failure.
+#define CNNL_CHECK(expr)                                                                                               \
+    do {                                                                                                               \
+        cnnlStatus_t _cnnl_status = (expr);                                                                            \
+        CHECK_EQ(_cnnl_status, CNNL_STATUS_SUCCESS) << "CNNL call failed: " #expr " -> " << _cnnl_status;              \
+    } while (0)
+
+} // namespace infini_train::common::mlu