diff --git a/backends/mlu/CMakeLists.txt b/backends/mlu/CMakeLists.txt new file mode 100644 index 0000000..eb4040b --- /dev/null +++ b/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 + "$" + 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() diff --git a/backends/mlu/cmake/pre_project.cmake b/backends/mlu/cmake/pre_project.cmake new file mode 100644 index 0000000..4e61bdf --- /dev/null +++ b/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=.") +endif() + +set(USE_CUDA OFF) +set(USE_NCCL OFF) +set(USE_OMP ON) +set(BUILD_SHARED_LIBS OFF) diff --git a/backends/mlu/examples/CMakeLists.txt b/backends/mlu/examples/CMakeLists.txt new file mode 100644 index 0000000..666e650 --- /dev/null +++ b/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) diff --git a/backends/mlu/include/infini_train_mlu/backend.h b/backends/mlu/include/infini_train_mlu/backend.h new file mode 100644 index 0000000..1bd968b --- /dev/null +++ b/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 diff --git a/backends/mlu/src/backend.cc b/backends/mlu/src/backend.cc new file mode 100644 index 0000000..cbbef62 --- /dev/null +++ b/backends/mlu/src/backend.cc @@ -0,0 +1,19 @@ +#include "infini_train_mlu/backend.h" + +#include + +#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 diff --git a/backends/mlu/src/common/common_mlu.h b/backends/mlu/src/common/common_mlu.h new file mode 100644 index 0000000..7fc22dd --- /dev/null +++ b/backends/mlu/src/common/common_mlu.h @@ -0,0 +1,24 @@ +#pragma once + +#include +#include + +#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 diff --git a/backends/mlu/src/kernels/cast.cc b/backends/mlu/src/kernels/cast.cc new file mode 100644 index 0000000..a50a4a4 --- /dev/null +++ b/backends/mlu/src/kernels/cast.cc @@ -0,0 +1,484 @@ +#include +#include + +#include + +#include "infini_train/include/core/runtime/device_guard.h" +#include "infini_train/include/device.h" +#include "infini_train/include/dispatcher.h" +#include "infini_train/include/tensor.h" + +#include "common/common_mlu.h" +#include "runtime/mlu_guard_impl.h" +#include "runtime/mlu_runtime_common.h" + +namespace infini_train::kernels::mlu { +namespace { + +cnnlDataType_t ToCnnlDType(DataType dtype) { + switch (dtype) { + case DataType::kBOOL: + return CNNL_DTYPE_BOOL; + case DataType::kUINT8: + return CNNL_DTYPE_UINT8; + case DataType::kINT8: + return CNNL_DTYPE_INT8; + case DataType::kUINT16: + return CNNL_DTYPE_UINT16; + case DataType::kINT16: + return CNNL_DTYPE_INT16; + case DataType::kUINT32: + return CNNL_DTYPE_UINT32; + case DataType::kINT32: + return CNNL_DTYPE_INT32; + case DataType::kUINT64: + return CNNL_DTYPE_UINT64; + case DataType::kINT64: + return CNNL_DTYPE_INT64; + case DataType::kBFLOAT16: + return CNNL_DTYPE_BFLOAT16; + case DataType::kFLOAT16: + return CNNL_DTYPE_HALF; + case DataType::kFLOAT32: + return CNNL_DTYPE_FLOAT; + case DataType::kFLOAT64: + return CNNL_DTYPE_DOUBLE; + default: + LOG(FATAL) << "MLU Cast: unsupported dtype " << static_cast(dtype); + } +} + +// Conversion selectors from the cnnlCastDataType_t table in cnnl.h. Returns +// false for pairs CNNL 2.2 does not provide (extend as the framework needs). +bool CastSelector(DataType src, DataType dst, cnnlCastDataType_t *selector) { + const auto s = static_cast(src); + const auto d = static_cast(dst); + switch (s) { + case static_cast(DataType::kFLOAT32): + switch (d) { + case static_cast(DataType::kFLOAT16): + *selector = CNNL_CAST_FLOAT_TO_HALF; + return true; + case static_cast(DataType::kBFLOAT16): + *selector = CNNL_CAST_FLOAT_TO_BFLOAT16; + return true; + case static_cast(DataType::kFLOAT64): + *selector = CNNL_CAST_FLOAT_TO_DOUBLE; + return true; + case static_cast(DataType::kINT32): + *selector = CNNL_CAST_FLOAT_TO_INT32; + return true; + case static_cast(DataType::kUINT32): + *selector = CNNL_CAST_FLOAT_TO_UINT32; + return true; + case static_cast(DataType::kINT16): + *selector = CNNL_CAST_FLOAT_TO_INT16; + return true; + case static_cast(DataType::kINT8): + *selector = CNNL_CAST_FLOAT_TO_INT8; + return true; + case static_cast(DataType::kUINT8): + *selector = CNNL_CAST_FLOAT_TO_UINT8; + return true; + case static_cast(DataType::kBOOL): + *selector = CNNL_CAST_FLOAT_TO_BOOL; + return true; + case static_cast(DataType::kINT64): + *selector = CNNL_CAST_FLOAT_TO_INT64; + return true; + case static_cast(DataType::kUINT64): + *selector = CNNL_CAST_FLOAT_TO_UINT64; + return true; + } + return false; + case static_cast(DataType::kFLOAT16): + switch (d) { + case static_cast(DataType::kFLOAT32): + *selector = CNNL_CAST_HALF_TO_FLOAT; + return true; + case static_cast(DataType::kBFLOAT16): + *selector = CNNL_CAST_HALF_TO_BFLOAT16; + return true; + case static_cast(DataType::kFLOAT64): + *selector = CNNL_CAST_HALF_TO_DOUBLE; + return true; + case static_cast(DataType::kINT32): + *selector = CNNL_CAST_HALF_TO_INT32; + return true; + case static_cast(DataType::kINT16): + *selector = CNNL_CAST_HALF_TO_INT16; + return true; + case static_cast(DataType::kINT8): + *selector = CNNL_CAST_HALF_TO_INT8; + return true; + case static_cast(DataType::kUINT8): + *selector = CNNL_CAST_HALF_TO_UINT8; + return true; + case static_cast(DataType::kUINT32): + *selector = CNNL_CAST_HALF_TO_UINT32; + return true; + case static_cast(DataType::kUINT64): + *selector = CNNL_CAST_HALF_TO_UINT64; + return true; + case static_cast(DataType::kINT64): + *selector = CNNL_CAST_HALF_TO_INT64; + return true; + case static_cast(DataType::kBOOL): + *selector = CNNL_CAST_HALF_TO_BOOL; + return true; + } + return false; + case static_cast(DataType::kBFLOAT16): + switch (d) { + case static_cast(DataType::kFLOAT32): + *selector = CNNL_CAST_BFLOAT16_TO_FLOAT; + return true; + case static_cast(DataType::kFLOAT16): + *selector = CNNL_CAST_BFLOAT16_TO_HALF; + return true; + case static_cast(DataType::kFLOAT64): + *selector = CNNL_CAST_BFLOAT16_TO_DOUBLE; + return true; + case static_cast(DataType::kINT8): + *selector = CNNL_CAST_BFLOAT16_TO_INT8; + return true; + case static_cast(DataType::kUINT8): + *selector = CNNL_CAST_BFLOAT16_TO_UINT8; + return true; + case static_cast(DataType::kINT16): + *selector = CNNL_CAST_BFLOAT16_TO_INT16; + return true; + case static_cast(DataType::kINT32): + *selector = CNNL_CAST_BFLOAT16_TO_INT32; + return true; + case static_cast(DataType::kUINT32): + *selector = CNNL_CAST_BFLOAT16_TO_UINT32; + return true; + case static_cast(DataType::kINT64): + *selector = CNNL_CAST_BFLOAT16_TO_INT64; + return true; + case static_cast(DataType::kUINT64): + *selector = CNNL_CAST_BFLOAT16_TO_UINT64; + return true; + case static_cast(DataType::kBOOL): + *selector = CNNL_CAST_BFLOAT16_TO_BOOL; + return true; + } + return false; + case static_cast(DataType::kFLOAT64): + switch (d) { + case static_cast(DataType::kFLOAT32): + *selector = CNNL_CAST_DOUBLE_TO_FLOAT; + return true; + case static_cast(DataType::kFLOAT16): + *selector = CNNL_CAST_DOUBLE_TO_HALF; + return true; + case static_cast(DataType::kBFLOAT16): + *selector = CNNL_CAST_DOUBLE_TO_BFLOAT16; + return true; + case static_cast(DataType::kINT64): + *selector = CNNL_CAST_DOUBLE_TO_INT64; + return true; + } + return false; + case static_cast(DataType::kINT32): + switch (d) { + case static_cast(DataType::kFLOAT32): + *selector = CNNL_CAST_INT32_TO_FLOAT; + return true; + case static_cast(DataType::kFLOAT16): + *selector = CNNL_CAST_INT32_TO_HALF; + return true; + case static_cast(DataType::kBFLOAT16): + *selector = CNNL_CAST_INT32_TO_BFLOAT16; + return true; + case static_cast(DataType::kINT8): + *selector = CNNL_CAST_INT32_TO_INT8; + return true; + case static_cast(DataType::kINT16): + *selector = CNNL_CAST_INT32_TO_INT16; + return true; + case static_cast(DataType::kUINT8): + *selector = CNNL_CAST_INT32_TO_UINT8; + return true; + case static_cast(DataType::kUINT32): + *selector = CNNL_CAST_INT32_TO_UINT32; + return true; + case static_cast(DataType::kUINT64): + *selector = CNNL_CAST_INT32_TO_UINT64; + return true; + case static_cast(DataType::kINT64): + *selector = CNNL_CAST_INT32_TO_INT64; + return true; + case static_cast(DataType::kBOOL): + *selector = CNNL_CAST_INT32_TO_BOOL; + return true; + } + return false; + case static_cast(DataType::kINT16): + switch (d) { + case static_cast(DataType::kFLOAT32): + *selector = CNNL_CAST_INT16_TO_FLOAT; + return true; + case static_cast(DataType::kFLOAT16): + *selector = CNNL_CAST_INT16_TO_HALF; + return true; + case static_cast(DataType::kBFLOAT16): + *selector = CNNL_CAST_INT16_TO_BFLOAT16; + return true; + case static_cast(DataType::kINT64): + *selector = CNNL_CAST_INT16_TO_INT64; + return true; + case static_cast(DataType::kUINT64): + *selector = CNNL_CAST_INT16_TO_UINT64; + return true; + case static_cast(DataType::kINT32): + *selector = CNNL_CAST_INT16_TO_INT32; + return true; + case static_cast(DataType::kUINT32): + *selector = CNNL_CAST_INT16_TO_UINT32; + return true; + case static_cast(DataType::kINT8): + *selector = CNNL_CAST_INT16_TO_INT8; + return true; + case static_cast(DataType::kBOOL): + *selector = CNNL_CAST_INT16_TO_BOOL; + return true; + case static_cast(DataType::kUINT8): + *selector = CNNL_CAST_INT16_TO_UINT8; + return true; + } + return false; + case static_cast(DataType::kINT8): + switch (d) { + case static_cast(DataType::kFLOAT32): + *selector = CNNL_CAST_INT8_TO_FLOAT; + return true; + case static_cast(DataType::kFLOAT16): + *selector = CNNL_CAST_INT8_TO_HALF; + return true; + case static_cast(DataType::kBFLOAT16): + *selector = CNNL_CAST_INT8_TO_BFLOAT16; + return true; + case static_cast(DataType::kUINT8): + *selector = CNNL_CAST_INT8_TO_UINT8; + return true; + case static_cast(DataType::kINT32): + *selector = CNNL_CAST_INT8_TO_INT32; + return true; + case static_cast(DataType::kUINT32): + *selector = CNNL_CAST_INT8_TO_UINT32; + return true; + case static_cast(DataType::kINT64): + *selector = CNNL_CAST_INT8_TO_INT64; + return true; + case static_cast(DataType::kUINT64): + *selector = CNNL_CAST_INT8_TO_UINT64; + return true; + case static_cast(DataType::kBOOL): + *selector = CNNL_CAST_INT8_TO_BOOL; + return true; + case static_cast(DataType::kINT16): + *selector = CNNL_CAST_INT8_TO_INT16; + return true; + } + return false; + case static_cast(DataType::kUINT8): + switch (d) { + case static_cast(DataType::kFLOAT32): + *selector = CNNL_CAST_UINT8_TO_FLOAT; + return true; + case static_cast(DataType::kFLOAT16): + *selector = CNNL_CAST_UINT8_TO_HALF; + return true; + case static_cast(DataType::kBFLOAT16): + *selector = CNNL_CAST_UINT8_TO_BFLOAT16; + return true; + case static_cast(DataType::kBOOL): + *selector = CNNL_CAST_UINT8_TO_BOOL; + return true; + case static_cast(DataType::kINT8): + *selector = CNNL_CAST_UINT8_TO_INT8; + return true; + case static_cast(DataType::kINT16): + *selector = CNNL_CAST_UINT8_TO_INT16; + return true; + case static_cast(DataType::kINT32): + *selector = CNNL_CAST_UINT8_TO_INT32; + return true; + case static_cast(DataType::kUINT32): + *selector = CNNL_CAST_UINT8_TO_UINT32; + return true; + case static_cast(DataType::kUINT64): + *selector = CNNL_CAST_UINT8_TO_UINT64; + return true; + case static_cast(DataType::kINT64): + *selector = CNNL_CAST_UINT8_TO_INT64; + return true; + } + return false; + case static_cast(DataType::kINT64): + switch (d) { + case static_cast(DataType::kINT32): + *selector = CNNL_CAST_INT64_TO_INT32; + return true; + case static_cast(DataType::kINT16): + *selector = CNNL_CAST_INT64_TO_INT16; + return true; + case static_cast(DataType::kINT8): + *selector = CNNL_CAST_INT64_TO_INT8; + return true; + case static_cast(DataType::kBOOL): + *selector = CNNL_CAST_INT64_TO_BOOL; + return true; + case static_cast(DataType::kFLOAT32): + *selector = CNNL_CAST_INT64_TO_FLOAT; + return true; + case static_cast(DataType::kFLOAT16): + *selector = CNNL_CAST_INT64_TO_HALF; + return true; + case static_cast(DataType::kBFLOAT16): + *selector = CNNL_CAST_INT64_TO_BFLOAT16; + return true; + case static_cast(DataType::kFLOAT64): + *selector = CNNL_CAST_INT64_TO_DOUBLE; + return true; + case static_cast(DataType::kUINT64): + *selector = CNNL_CAST_INT64_TO_UINT64; + return true; + case static_cast(DataType::kUINT32): + *selector = CNNL_CAST_INT64_TO_UINT32; + return true; + case static_cast(DataType::kUINT8): + *selector = CNNL_CAST_INT64_TO_UINT8; + return true; + } + return false; + case static_cast(DataType::kUINT64): + switch (d) { + case static_cast(DataType::kFLOAT16): + *selector = CNNL_CAST_UINT64_TO_HALF; + return true; + case static_cast(DataType::kBFLOAT16): + *selector = CNNL_CAST_UINT64_TO_BFLOAT16; + return true; + case static_cast(DataType::kFLOAT32): + *selector = CNNL_CAST_UINT64_TO_FLOAT; + return true; + case static_cast(DataType::kUINT8): + *selector = CNNL_CAST_UINT64_TO_UINT8; + return true; + case static_cast(DataType::kBOOL): + *selector = CNNL_CAST_UINT64_TO_BOOL; + return true; + case static_cast(DataType::kUINT32): + *selector = CNNL_CAST_UINT64_TO_UINT32; + return true; + } + return false; + case static_cast(DataType::kUINT32): + switch (d) { + case static_cast(DataType::kFLOAT32): + *selector = CNNL_CAST_UINT32_TO_FLOAT; + return true; + case static_cast(DataType::kFLOAT16): + *selector = CNNL_CAST_UINT32_TO_HALF; + return true; + case static_cast(DataType::kBFLOAT16): + *selector = CNNL_CAST_UINT32_TO_BFLOAT16; + return true; + case static_cast(DataType::kUINT8): + *selector = CNNL_CAST_UINT32_TO_UINT8; + return true; + case static_cast(DataType::kINT64): + *selector = CNNL_CAST_UINT32_TO_INT64; + return true; + case static_cast(DataType::kUINT64): + *selector = CNNL_CAST_UINT32_TO_UINT64; + return true; + case static_cast(DataType::kBOOL): + *selector = CNNL_CAST_UINT32_TO_BOOL; + return true; + } + return false; + case static_cast(DataType::kBOOL): + switch (d) { + case static_cast(DataType::kFLOAT32): + *selector = CNNL_CAST_BOOL_TO_FLOAT; + return true; + case static_cast(DataType::kFLOAT16): + *selector = CNNL_CAST_BOOL_TO_HALF; + return true; + case static_cast(DataType::kBFLOAT16): + *selector = CNNL_CAST_BOOL_TO_BFLOAT16; + return true; + case static_cast(DataType::kINT8): + *selector = CNNL_CAST_BOOL_TO_INT8; + return true; + case static_cast(DataType::kUINT8): + *selector = CNNL_CAST_BOOL_TO_UINT8; + return true; + case static_cast(DataType::kINT16): + *selector = CNNL_CAST_BOOL_TO_INT16; + return true; + case static_cast(DataType::kINT32): + *selector = CNNL_CAST_BOOL_TO_INT32; + return true; + case static_cast(DataType::kINT64): + *selector = CNNL_CAST_BOOL_TO_INT64; + return true; + case static_cast(DataType::kUINT32): + *selector = CNNL_CAST_BOOL_TO_UINT32; + return true; + case static_cast(DataType::kUINT64): + *selector = CNNL_CAST_BOOL_TO_UINT64; + return true; + } + return false; + default: + return false; + } +} +} // namespace + +std::shared_ptr Cast(std::shared_ptr input, DataType dtype) { + auto dst_tensor = std::make_shared(input->Dims(), dtype, input->GetDevice()); + auto device = input->GetDevice(); + auto *stream = core::GetDeviceGuardImpl(device.type())->GetStream(device); + const auto &mlu_queue = dynamic_cast(stream)->mlu_queue(); + core::mlu::MluBlasHandle blas_handle(stream); + + if (input->Dtype() == dtype) { // identity cast: plain device-to-device copy + core::GetDeviceGuardImpl(device.type()) + ->Memcpy(dst_tensor->DataPtr(), input->DataPtr(), input->NumElements() * kDataTypeToSize.at(dtype), + core::MemcpyKind::kD2D); + return dst_tensor; + } + + cnnlCastDataType_t selector = CNNL_CAST_FLOAT_TO_HALF; + CHECK(CastSelector(input->Dtype(), dtype, &selector)) + << "MLU Cast: no CNNL conversion from " << static_cast(input->Dtype()) << " to " + << static_cast(dtype); + + cnnlTensorDescriptor_t x_desc = nullptr; + cnnlTensorDescriptor_t y_desc = nullptr; + CNNL_CHECK(cnnlCreateTensorDescriptor(&x_desc)); + CNNL_CHECK(cnnlCreateTensorDescriptor(&y_desc)); + CHECK_LE(input->NumElements(), 0x7fffffffull) << "MLU Cast: tensor too large"; + int dims[1] = {static_cast(input->NumElements())}; + CNNL_CHECK(cnnlSetTensorDescriptor(x_desc, CNNL_LAYOUT_ARRAY, ToCnnlDType(input->Dtype()), 1, dims)); + CNNL_CHECK(cnnlSetTensorDescriptor(y_desc, CNNL_LAYOUT_ARRAY, ToCnnlDType(dtype), 1, dims)); + CNNL_CHECK( + cnnlCastDataType(blas_handle.cnnl_handle(), x_desc, input->DataPtr(), selector, y_desc, dst_tensor->DataPtr())); + CNNL_CHECK(cnnlDestroyTensorDescriptor(x_desc)); + CNNL_CHECK(cnnlDestroyTensorDescriptor(y_desc)); + + return dst_tensor; +} +} // namespace infini_train::kernels::mlu + +namespace infini_train::kernels::mlu { + +REGISTER_KERNEL(infini_train::Device::DeviceType::kPrivateUse1, Cast, infini_train::kernels::mlu::Cast); + +} // namespace infini_train::kernels::mlu diff --git a/backends/mlu/src/kernels/fill.cc b/backends/mlu/src/kernels/fill.cc new file mode 100644 index 0000000..47dfa24 --- /dev/null +++ b/backends/mlu/src/kernels/fill.cc @@ -0,0 +1,117 @@ +#include +#include +#include + +#include "infini_train/include/core/runtime/device_guard.h" +#include "infini_train/include/device.h" +#include "infini_train/include/dispatcher.h" +#include "infini_train/include/scalar.h" +#include "infini_train/include/tensor.h" + +#include "runtime/mlu_guard_impl.h" +#include "runtime/mlu_runtime_common.h" + +namespace infini_train::kernels::mlu { +namespace { + +extern "C" void infini_train_mlu_fill_pattern(void *dst, unsigned int elemsize, unsigned long long pattern, + unsigned int num_elements, cnrtDim3_t dim, cnrtFunctionType_t func_type, + cnrtQueue_t queue); + +// Element bit pattern for the fill constant, converted with the same +// Scalar::to() semantics the CPU kernel uses. +unsigned long long BuildPattern(const Scalar &scalar, DataType dtype) { + unsigned long long pattern = 0; + switch (dtype) { + case DataType::kBOOL: { + unsigned char v = scalar.to() ? 1 : 0; + std::memcpy(&pattern, &v, sizeof(v)); + break; + } + case DataType::kUINT8: { + auto v = scalar.to(); + std::memcpy(&pattern, &v, sizeof(v)); + break; + } + case DataType::kINT8: { + auto v = scalar.to(); + std::memcpy(&pattern, &v, sizeof(v)); + break; + } + case DataType::kUINT16: { + auto v = scalar.to(); + std::memcpy(&pattern, &v, sizeof(v)); + break; + } + case DataType::kINT16: { + auto v = scalar.to(); + std::memcpy(&pattern, &v, sizeof(v)); + break; + } + case DataType::kUINT32: { + auto v = scalar.to(); + std::memcpy(&pattern, &v, sizeof(v)); + break; + } + case DataType::kINT32: { + auto v = scalar.to(); + std::memcpy(&pattern, &v, sizeof(v)); + break; + } + case DataType::kUINT64: { + auto v = scalar.to(); + std::memcpy(&pattern, &v, sizeof(v)); + break; + } + case DataType::kINT64: { + auto v = scalar.to(); + std::memcpy(&pattern, &v, sizeof(v)); + break; + } + case DataType::kBFLOAT16: { + BF16 v = scalar.to(); + std::memcpy(&pattern, &v, sizeof(v)); + break; + } + case DataType::kFLOAT16: { + FP16 v = scalar.to(); + std::memcpy(&pattern, &v, sizeof(v)); + break; + } + case DataType::kFLOAT32: { + float v = scalar.to(); + std::memcpy(&pattern, &v, sizeof(v)); + break; + } + case DataType::kFLOAT64: { + double v = scalar.to(); + std::memcpy(&pattern, &v, sizeof(v)); + break; + } + default: + LOG(FATAL) << "MLU Fill: unsupported dtype " << static_cast(dtype); + } + return pattern; +} +} // namespace + +void Fill(std::shared_ptr tensor, Scalar scalar) { + auto device = tensor->GetDevice(); + auto *stream = core::GetDeviceGuardImpl(device.type())->GetStream(device); + const auto &mlu_stream = dynamic_cast(stream)->mlu_queue(); + + const size_t elemsize = kDataTypeToSize.at(tensor->Dtype()); + CHECK_LE(tensor->NumElements(), 0xffffffffull) << "MLU Fill: tensor too large"; + + unsigned long long pattern = BuildPattern(scalar, tensor->Dtype()); + cnrtDim3_t dim = {4, 1, 1}; + infini_train_mlu_fill_pattern(tensor->DataPtr(), (unsigned int)elemsize, pattern, + (unsigned int)tensor->NumElements(), dim, cnrtFuncTypeUnion1, mlu_stream); +} +} // namespace infini_train::kernels::mlu + +namespace infini_train::kernels::mlu { + +REGISTER_KERNEL(infini_train::Device::DeviceType::kPrivateUse1, Fill, infini_train::kernels::mlu::Fill); + +} // namespace infini_train::kernels::mlu diff --git a/backends/mlu/src/kernels/fill.mlu b/backends/mlu/src/kernels/fill.mlu new file mode 100644 index 0000000..809377a --- /dev/null +++ b/backends/mlu/src/kernels/fill.mlu @@ -0,0 +1,73 @@ +// BANG device kernels for the Fill op, plus the dtype-agnostic host launcher. +// +// Fill fills every element of a device buffer with a constant. The value +// travels from the host as an element-sized bit pattern, so a single pair of +// kernels covers every framework dtype: 8-byte and 4-byte element buffers are +// filled with the pattern word directly; 2-byte and 1-byte element buffers are +// filled through a uniform 4-byte repeat plus a small byte-wise tail. + +#include "bang.h" +#include "cnrt.h" + +#include + +namespace { + +template __mlu_global__ void FillPatternKernel(T *dst, T pattern, unsigned int n) { + // Contiguous per-task chunking: taskId in [0, taskDim), one core per task. + unsigned int per = (n + taskDim - 1) / taskDim; + unsigned int begin = taskId * per; + unsigned int end = begin + per; + if (end > n) { + end = n; + } + for (unsigned int i = begin; i < end; i++) { dst[i] = pattern; } +} + +__mlu_global__ void FillTailKernel(unsigned char *dst, unsigned int pattern, unsigned int n) { + unsigned int per = (n + taskDim - 1) / taskDim; + unsigned int begin = taskId * per; + unsigned int end = begin + per; + if (end > n) { + end = n; + } + for (unsigned int i = begin; i < end; i++) { dst[i] = (unsigned char)((pattern >> (8 * i)) & 0xffu); } +} + +} // namespace + +// dtype_tag is the element size in bytes (1/2/4/8); pattern carries the +// element bit pattern in its low bytes. num_elements fits in 32 bits for any +// tensor the framework can allocate today (checked by the host side). +extern "C" void infini_train_mlu_fill_pattern(void *dst, unsigned int elemsize, unsigned long long pattern, + unsigned int num_elements, cnrtDim3_t dim, cnrtFunctionType_t func_type, + cnrtQueue_t queue) { + if (elemsize == 8) { + FillPatternKernel + <<>>((unsigned long long *)dst, pattern, num_elements); + return; + } + if (elemsize == 4) { + FillPatternKernel + <<>>((unsigned int *)dst, (unsigned int)pattern, num_elements); + return; + } + unsigned char *base = (unsigned char *)dst; + unsigned int total_bytes = num_elements * elemsize; + unsigned int pat32; + if (elemsize == 2) { + unsigned int lo16 = (unsigned int)(pattern & 0xffffu); + pat32 = lo16 | (lo16 << 16); + } else { + unsigned int byte = (unsigned int)(pattern & 0xffu); + pat32 = byte | (byte << 8) | (byte << 16) | (byte << 24); + } + unsigned int n32 = total_bytes / 4; + if (n32 > 0) { + FillPatternKernel<<>>((unsigned int *)dst, pat32, n32); + } + unsigned int tail_bytes = total_bytes - n32 * 4; + if (tail_bytes > 0) { + FillTailKernel<<>>(base + n32 * 4, pattern, tail_bytes); + } +} diff --git a/backends/mlu/src/kernels/no_op.cc b/backends/mlu/src/kernels/no_op.cc new file mode 100644 index 0000000..57fba3e --- /dev/null +++ b/backends/mlu/src/kernels/no_op.cc @@ -0,0 +1,38 @@ +#include +#include +#include +#include +#include + +#include "glog/logging.h" + +#include "infini_train/include/dispatcher.h" +#include "infini_train/include/tensor.h" + +namespace infini_train::kernels::mlu { + +// Pure host-side views (Tensor(Tensor, offset, dims) aliases the storage), so +// the MLU implementation is identical to the CPU/MACA ones - no device work. +std::shared_ptr NoOpForward(const std::shared_ptr &input, const std::vector &dims) { + const int64_t num_elements = std::accumulate(dims.begin(), dims.end(), 1, std::multiplies()); + CHECK_EQ(input->NumElements(), num_elements); + + auto output = std::make_shared(*input, 0, dims); + return output; +} + +std::shared_ptr NoOpBackward(const std::vector &dims, const std::shared_ptr &grad_output) { + auto num_elements = std::accumulate(dims.begin(), dims.end(), 1, std::multiplies()); + CHECK_EQ(num_elements, grad_output->NumElements()); + + auto grad_input = std::make_shared(*grad_output, 0, dims); + return grad_input; +} +} // namespace infini_train::kernels::mlu + +namespace infini_train::kernels::mlu { + +REGISTER_KERNEL(infini_train::Device::DeviceType::kPrivateUse1, NoOpForward, infini_train::kernels::mlu::NoOpForward) +REGISTER_KERNEL(infini_train::Device::DeviceType::kPrivateUse1, NoOpBackward, infini_train::kernels::mlu::NoOpBackward) + +} // namespace infini_train::kernels::mlu diff --git a/backends/mlu/src/runtime/mlu_guard_impl.cc b/backends/mlu/src/runtime/mlu_guard_impl.cc new file mode 100644 index 0000000..cc543a6 --- /dev/null +++ b/backends/mlu/src/runtime/mlu_guard_impl.cc @@ -0,0 +1,378 @@ +#include "runtime/mlu_guard_impl.h" + +#include +#include +#include +#include +#include +#include +#include +#include + +#include "glog/logging.h" + +#include "infini_train/include/core/runtime/runtime_common.h" +#include "infini_train/include/device.h" + +#include "common/common_mlu.h" +#include "runtime/mlu_runtime_common.h" + +namespace infini_train::core::mlu { +namespace { +constexpr size_t kBytesPerMB = 1024ULL * 1024ULL; + +static std::vector> g_mlu_streams; +static std::vector> g_mlu_blas_handles; +static std::vector> g_device_stream_flags; +static std::vector> g_device_handle_flags; + +struct AllocationRecord { + int device_index = -1; + size_t size = 0; +}; + +struct AllocationStats { + size_t current = 0; + size_t high = 0; +}; + +// CNRT allocations go straight to the driver with no caching pool, so the +// allocation tracker below is the only source of memory statistics. +static std::mutex g_allocation_mutex; +static std::unordered_map g_allocations; +static std::vector g_allocation_stats; + +void TrackAllocation(void *ptr, int device_index, size_t size) { + if (ptr == nullptr) { + return; + } + + std::lock_guard lock(g_allocation_mutex); + g_allocations[ptr] = AllocationRecord{device_index, size}; + auto &stats = g_allocation_stats.at(device_index); + stats.current += size; + stats.high = std::max(stats.high, stats.current); +} + +void UntrackAllocation(void *ptr) { + if (ptr == nullptr) { + return; + } + + std::lock_guard lock(g_allocation_mutex); + auto it = g_allocations.find(ptr); + if (it == g_allocations.end()) { + return; + } + + auto &stats = g_allocation_stats.at(it->second.device_index); + CHECK_GE(stats.current, it->second.size) << "MLU allocation tracker underflow."; + stats.current -= it->second.size; + g_allocations.erase(it); +} + +inline void CheckMluDevice(Device device) { + CHECK(device.type() == Device::DeviceType::kPrivateUse1) + << "MluGuardImpl expects MLU device, but got type=" << static_cast(device.type()) + << " index=" << device.index(); + const int idx = device.index(); + CHECK(idx >= 0 && static_cast(idx) < g_mlu_streams.size()) + << "MLU device index " << idx << " out of cache range [0, " << g_mlu_streams.size() << ")."; +} + +inline cnrtNotifier_t GetMluEvent(Event *event) { + auto *mlu_event = dynamic_cast(event); + CHECK_NOTNULL(mlu_event); + return mlu_event->mlu_event(); +} + +inline cnrtQueue_t GetMluQueue(Stream *stream) { + auto *mlu_stream = dynamic_cast(stream); + CHECK_NOTNULL(mlu_stream); + return mlu_stream->mlu_queue(); +} + +cnrtMemTransDir_t ToCnrtMemcpyKind(MemcpyKind kind, const char *context) { + switch (kind) { + case MemcpyKind::kH2D: + return cnrtMemcpyHostToDev; + case MemcpyKind::kD2H: + return cnrtMemcpyDevToHost; + case MemcpyKind::kD2D: + return cnrtMemcpyDevToDev; + default: + LOG(FATAL) << context << " got invalid MemcpyKind=" << MemcpyKindToString(kind); + } +} +} // namespace + +void MluGuardImpl::InitSingleStream(Device device) { + CheckMluDevice(device); + + int current_device = -1; + MLU_CHECK(cnrtGetDevice(¤t_device)); + MLU_CHECK(cnrtSetDevice(device.index())); + + g_mlu_streams[device.index()] = std::make_unique(); + + MLU_CHECK(cnrtSetDevice(current_device)); +} + +void MluGuardImpl::InitSingleHandle(Device device) { + CheckMluDevice(device); + + int current_device = -1; + MLU_CHECK(cnrtGetDevice(¤t_device)); + MLU_CHECK(cnrtSetDevice(device.index())); + + std::call_once(*g_device_stream_flags.at(device.index()), InitSingleStream, device); + + g_mlu_blas_handles[device.index()] = std::make_unique(g_mlu_streams[device.index()].get()); + + MLU_CHECK(cnrtSetDevice(current_device)); +} + +void MluGuardImpl::Initialize() { + std::call_once(initialize_flag_, [] { + // CNRT 7.x initializes lazily on the first device API call; there is no + // explicit global init entry point to invoke here. + unsigned int device_count = 0; + MLU_CHECK(cnrtGetDeviceCount(&device_count)); + CHECK_GT(static_cast(device_count), 0) << "No MLU devices are available."; + CHECK_LE(static_cast(device_count), static_cast(std::numeric_limits::max()) + 1) + << "MLU device count exceeds InfiniTrain's Device index range."; + + g_mlu_streams.resize(device_count); + g_mlu_blas_handles.resize(device_count); + g_allocation_stats.resize(device_count); + g_device_stream_flags.reserve(device_count); + g_device_handle_flags.reserve(device_count); + for (unsigned int i = 0; i < device_count; ++i) { + g_device_stream_flags.push_back(std::make_unique()); + g_device_handle_flags.push_back(std::make_unique()); + } + }); +} + +// device +Device MluGuardImpl::GetDevice() const { + int current_device = -1; + MLU_CHECK(cnrtGetDevice(¤t_device)); + return Device(Device::DeviceType::kPrivateUse1, current_device); +} + +void MluGuardImpl::SetDevice(Device device) const { + CheckMluDevice(device); + MLU_CHECK(cnrtSetDevice(device.index())); +} + +int MluGuardImpl::DeviceCount() const { return static_cast(g_mlu_streams.size()); } + +Device::DeviceType MluGuardImpl::Type() const { return Device::DeviceType::kPrivateUse1; } + +// stream +Stream *MluGuardImpl::GetStream(Device device) const { + CheckMluDevice(device); + std::call_once(*g_device_stream_flags.at(device.index()), InitSingleStream, device); + return g_mlu_streams.at(device.index()).get(); +} + +Stream *MluGuardImpl::CreateStream(Device device) const { + CheckMluDevice(device); + int current_device = -1; + MLU_CHECK(cnrtGetDevice(¤t_device)); + MLU_CHECK(cnrtSetDevice(device.index())); + + Stream *stream = new MluStream(); + + MLU_CHECK(cnrtSetDevice(current_device)); + return stream; +} + +Stream *MluGuardImpl::CreateStreamWithPriority(Device device, int priority) const { + CheckMluDevice(device); + int current_device = -1; + MLU_CHECK(cnrtGetDevice(¤t_device)); + MLU_CHECK(cnrtSetDevice(device.index())); + + Stream *stream = new MluStream(priority); + + MLU_CHECK(cnrtSetDevice(current_device)); + return stream; +} + +void MluGuardImpl::DestroyStream(Stream *stream) const { + if (stream == nullptr) { + return; + } + auto *mlu_stream = dynamic_cast(stream); + CHECK_NOTNULL(mlu_stream); + MLU_CHECK(cnrtQueueDestroy(mlu_stream->mlu_queue())); + delete mlu_stream; +} + +void MluGuardImpl::GetStreamPriorityRange(int *low, int *high) const { + // CNRT reports (min, max) accepted by cnrtQueueCreateWithPriority; pass the + // range through in the framework's (low, high) convention. + MLU_CHECK(cnrtDeviceGetQueuePriorityRange(low, high)); +} + +// event +void MluGuardImpl::EventCreate(Event **event) const { *event = new MluEvent(); } + +void MluGuardImpl::EventCreateWithFlags(Event **event, EventFlag flags) const { *event = new MluEvent(flags); } + +void MluGuardImpl::EventDestroy(Event *event) const { + if (event == nullptr) { + return; + } + // The wrapped notifier is intentionally not destroyed here; see the + // MluStream/MluEvent teardown notes. + delete event; +} + +void MluGuardImpl::EventRecord(Event *event, Stream *stream) const { + auto mlu_event = GetMluEvent(event); + auto mlu_queue = GetMluQueue(stream); + MLU_CHECK(cnrtPlaceNotifier(mlu_event, mlu_queue)); +} + +void MluGuardImpl::StreamWaitEvent(Stream *stream, Event *event, uint32_t flags) const { + auto mlu_event = GetMluEvent(event); + auto mlu_queue = GetMluQueue(stream); + MLU_CHECK(cnrtQueueWaitNotifier(mlu_event, mlu_queue, flags)); +} + +RuntimeStatus MluGuardImpl::EventSynchronize(Event *event) const { + auto mlu_event = GetMluEvent(event); + cnrtRet_t status = cnrtWaitNotifier(mlu_event); + if (status == cnrtSuccess) { + return RuntimeStatus::kSuccess; + } + if (status == cnrtErrorNotReady) { + return RuntimeStatus::kNotReady; + } + LOG(ERROR) << "MluGuardImpl::EventSynchronize failed: " << cnrtGetErrorStr(status); + return RuntimeStatus::kError; +} + +RuntimeStatus MluGuardImpl::EventQuery(Event *event) const { + auto mlu_event = GetMluEvent(event); + cnrtRet_t status = cnrtQueryNotifier(mlu_event); + if (status == cnrtSuccess) { + return RuntimeStatus::kSuccess; + } + if (status == cnrtErrorNotReady) { + return RuntimeStatus::kNotReady; + } + LOG(ERROR) << "MluGuardImpl::EventQuery failed: " << cnrtGetErrorStr(status); + return RuntimeStatus::kError; +} + +float MluGuardImpl::EventElapsedTime(Event *start_event, Event *stop_event) const { + auto start_mlu_event = GetMluEvent(start_event); + auto stop_mlu_event = GetMluEvent(stop_event); + float elapsed_ms = 0.0f; + MLU_CHECK(cnrtNotifierDuration(start_mlu_event, stop_mlu_event, &elapsed_ms)); + return elapsed_ms; +} + +// sync +void MluGuardImpl::SynchronizeDevice(Device device) const { + auto original_device = GetDevice(); + SetDevice(device); + + MLU_CHECK(cnrtSyncDevice()); + + SetDevice(original_device); +} + +void MluGuardImpl::SynchronizeStream(Stream *stream) const { + auto mlu_queue = GetMluQueue(stream); + MLU_CHECK(cnrtQueueSync(mlu_queue)); +} + +// blas +BlasHandle *MluGuardImpl::GetBlasHandle(Device device) const { + CheckMluDevice(device); + std::call_once(*g_device_handle_flags.at(device.index()), InitSingleHandle, device); + return g_mlu_blas_handles.at(device.index()).get(); +} + +// memory +void MluGuardImpl::Malloc(void **dev_ptr, size_t size) { + const cnrtRet_t status = cnrtMalloc(dev_ptr, size); + if (status != cnrtSuccess) { + int device_index = -1; + size_t free_bytes = 0, total_bytes = 0; + const cnrtRet_t device_status = cnrtGetDevice(&device_index); + const cnrtRet_t memory_status = cnrtMemGetInfo(&free_bytes, &total_bytes); + LOG(ERROR) << "MLU allocation failed: requested_bytes=" << size << " device=" << device_index + << " free_bytes=" << free_bytes << " total_bytes=" << total_bytes + << " device_query=" << cnrtGetErrorStr(device_status) + << " memory_query=" << cnrtGetErrorStr(memory_status); + } + MLU_CHECK(status); + + int device_index = -1; + MLU_CHECK(cnrtGetDevice(&device_index)); + TrackAllocation(*dev_ptr, device_index, size); +} + +void MluGuardImpl::MallocAsync(void **dev_ptr, size_t size, Stream *stream) { + // CNRT 7.x has no asynchronous allocation API; allocate synchronously on + // the allocation device. The stream argument is unused. + (void)stream; + Malloc(dev_ptr, size); +} + +void MluGuardImpl::Free(void *dev_ptr) { + UntrackAllocation(dev_ptr); + MLU_CHECK(cnrtFree(dev_ptr)); +} + +void MluGuardImpl::FreeAsync(void *dev_ptr, Stream *stream) { + if (dev_ptr == nullptr) { + return; + } + // With synchronous allocation there may still be queued work touching this + // block, so wait for the stream before releasing it. + SynchronizeStream(stream); + Free(dev_ptr); +} + +void MluGuardImpl::Memcpy(void *dst, const void *src, size_t count, MemcpyKind kind) { + const cnrtMemTransDir_t cnrt_kind = ToCnrtMemcpyKind(kind, "MluGuardImpl::Memcpy"); + MLU_CHECK(cnrtMemcpy(dst, const_cast(src), count, cnrt_kind)); +} + +void MluGuardImpl::MemcpyAsync(void *dst, const void *src, size_t count, MemcpyKind kind, Stream *stream) { + const cnrtMemTransDir_t cnrt_kind = ToCnrtMemcpyKind(kind, "MluGuardImpl::MemcpyAsync"); + auto mlu_queue = GetMluQueue(stream); + MLU_CHECK(cnrtMemcpyAsync(dst, const_cast(src), count, mlu_queue, cnrt_kind)); +} + +void MluGuardImpl::ResetMemPoolHighWatermarks(Device device) const { + CheckMluDevice(device); + + std::lock_guard lock(g_allocation_mutex); + auto &stats = g_allocation_stats.at(device.index()); + stats.high = stats.current; +} + +std::pair MluGuardImpl::GetMemPoolPeakMB(Device device) const { + CheckMluDevice(device); + + size_t tracked_high = 0; + { + std::lock_guard lock(g_allocation_mutex); + tracked_high = g_allocation_stats.at(device.index()).high; + } + + const size_t peak_mb = tracked_high / kBytesPerMB; + return {peak_mb, peak_mb}; +} + +INFINI_TRAIN_REGISTER_DEVICE_GUARD_IMPL(Device::DeviceType::kPrivateUse1, MluGuardImpl) + +} // namespace infini_train::core::mlu diff --git a/backends/mlu/src/runtime/mlu_guard_impl.h b/backends/mlu/src/runtime/mlu_guard_impl.h new file mode 100644 index 0000000..8a374d5 --- /dev/null +++ b/backends/mlu/src/runtime/mlu_guard_impl.h @@ -0,0 +1,92 @@ +#pragma once + +#include +#include +#include +#include + +#include "infini_train/include/core/runtime/device_guard.h" + +namespace infini_train::core { +class Stream; +class BlasHandle; +} // namespace infini_train::core + +namespace infini_train::core::mlu { + +class MluGuardImpl final : public DeviceGuardImpl { +public: + static void InitSingleStream(Device device); + + static void InitSingleHandle(Device device); + + MluGuardImpl() = default; + void Initialize() override; + + // device + Device GetDevice() const override; + + void SetDevice(Device device) const override; + + int DeviceCount() const override; + + Device::DeviceType Type() const override; + + // stream + Stream *GetStream(Device device) const override; + + Stream *CreateStream(Device device) const override; + + Stream *CreateStreamWithPriority(Device device, int priority) const override; + + void DestroyStream(Stream *stream) const override; + + void GetStreamPriorityRange(int *low, int *high) const override; + + // event + void EventCreate(Event **event) const override; + + void EventCreateWithFlags(Event **event, EventFlag flags) const override; + + void EventDestroy(Event *event) const override; + + void EventRecord(Event *event, Stream *stream) const override; + + void StreamWaitEvent(Stream *stream, Event *event, uint32_t flags) const override; + + RuntimeStatus EventSynchronize(Event *event) const override; + + RuntimeStatus EventQuery(Event *event) const override; + + float EventElapsedTime(Event *start_event, Event *stop_event) const override; + + // sync + void SynchronizeDevice(Device device) const override; + + void SynchronizeStream(Stream *stream) const override; + + // blas + BlasHandle *GetBlasHandle(Device device) const override; + + // memory + void Malloc(void **dev_ptr, size_t size) override; + + void MallocAsync(void **dev_ptr, size_t size, Stream *stream) override; + + void Free(void *dev_ptr) override; + + void FreeAsync(void *dev_ptr, Stream *stream) override; + + void Memcpy(void *dst, const void *src, size_t count, MemcpyKind kind) override; + + void MemcpyAsync(void *dst, const void *src, size_t count, MemcpyKind kind, Stream *stream) override; + + void ResetMemPoolHighWatermarks(Device device) const override; + + std::pair GetMemPoolPeakMB(Device device) const override; + +private: + std::once_flag initialize_flag_; +}; + +} // namespace infini_train::core::mlu diff --git a/backends/mlu/src/runtime/mlu_runtime_common.cc b/backends/mlu/src/runtime/mlu_runtime_common.cc new file mode 100644 index 0000000..fd708d7 --- /dev/null +++ b/backends/mlu/src/runtime/mlu_runtime_common.cc @@ -0,0 +1,52 @@ +#include "runtime/mlu_runtime_common.h" + +#include "common/common_mlu.h" + +namespace infini_train::core::mlu { + +unsigned int ToCnrtNotifierFlags(EventFlag flags) { + switch (flags) { + case EventFlag::kBlockingSync: + return CNRT_NOTIFIER_SYNC_WAIT; + case EventFlag::kDisableTiming: + return CNRT_NOTIFIER_DISABLE_TIMING; + case EventFlag::kInterprocess: + LOG(WARNING) << "MLU notifiers do not support interprocess events; falling back to default flags."; + return CNRT_NOTIFIER_DEFAULT; + case EventFlag::kDefault: + return CNRT_NOTIFIER_DEFAULT; + } + LOG(FATAL) << "Unsupported EventFlag: " << static_cast(flags); + return CNRT_NOTIFIER_DEFAULT; +} + +MluEvent::MluEvent(EventFlag flags) { MLU_CHECK(cnrtNotifierCreateWithFlags(&event_, ToCnrtNotifierFlags(flags))); } + +// See the class comment: no CNRT calls during teardown. +MluEvent::~MluEvent() = default; + +cnrtNotifier_t MluEvent::mlu_event() const { return event_; } + +MluStream::MluStream() { MLU_CHECK(cnrtQueueCreate(&queue_)); } + +MluStream::MluStream(int priority) { MLU_CHECK(cnrtQueueCreateWithPriority(&queue_, 0, priority)); } + +// See the class comment: no CNRT calls during teardown. +MluStream::~MluStream() = default; + +cnrtQueue_t MluStream::mlu_queue() const { return queue_; } + +MluBlasHandle::MluBlasHandle(Stream *stream) { + auto *mlu_stream = dynamic_cast(stream); + CHECK_NOTNULL(mlu_stream); + + CNNL_CHECK(cnnlCreate(&cnnl_handle_)); + CNNL_CHECK(cnnlSetQueue(cnnl_handle_, mlu_stream->mlu_queue())); +} + +// See the class comment: no CNNL calls during teardown. +MluBlasHandle::~MluBlasHandle() = default; + +cnnlHandle_t MluBlasHandle::cnnl_handle() const { return cnnl_handle_; } + +} // namespace infini_train::core::mlu diff --git a/backends/mlu/src/runtime/mlu_runtime_common.h b/backends/mlu/src/runtime/mlu_runtime_common.h new file mode 100644 index 0000000..1c8ab55 --- /dev/null +++ b/backends/mlu/src/runtime/mlu_runtime_common.h @@ -0,0 +1,60 @@ +#pragma once + +#include +#include + +#include "infini_train/include/core/runtime/runtime_common.h" + +namespace infini_train::core { +class Stream; +} + +namespace infini_train::core::mlu { + +// Maps a framework EventFlag onto CNRT notifier creation flags. +unsigned int ToCnrtNotifierFlags(EventFlag flags); + +class MluEvent final : public Event { +public: + explicit MluEvent(EventFlag flags = EventFlag::kDefault); + ~MluEvent() override; + + cnrtNotifier_t mlu_event() const; + +private: + cnrtNotifier_t event_ = nullptr; +}; + +class MluStream final : public Stream { +public: + MluStream(); + explicit MluStream(int priority); + + // NOTE: + // The wrapper destructor avoids CNRT calls because global destruction may + // outlive the runtime. Explicitly owned streams must be released through + // DeviceGuardImpl::DestroyStream while the runtime is alive. + ~MluStream() override; + + cnrtQueue_t mlu_queue() const; + +private: + cnrtQueue_t queue_ = nullptr; +}; + +class MluBlasHandle final : public BlasHandle { +public: + explicit MluBlasHandle(Stream *stream); + + // NOTE: + // Mirror MluStream: leaked intentionally; released only when the runtime + // is guaranteed to still be alive. + ~MluBlasHandle() override; + + cnnlHandle_t cnnl_handle() const; + +private: + cnnlHandle_t cnnl_handle_ = nullptr; +}; + +} // namespace infini_train::core::mlu