From 253263576cfcc3dc7f6fdcdf32d37e561ab0e45b Mon Sep 17 00:00:00 2001 From: cy <2833839179@qq.com> Date: Wed, 5 Aug 2026 10:28:58 +0800 Subject: [PATCH 1/4] feat(dcu): add BW1000 DCU backend support --- CMakeLists.txt | 166 +++ README.md | 13 +- example/gpt2/checkpoint_loader.cc | 1 + example/gpt2/main.cc | 36 +- example/llama3/checkpoint_loader.cc | 1 + example/llama3/main.cc | 36 +- format | 147 ++ infini_train/include/common/dcu/common_dcu.h | 46 + .../include/common/dcu/cub_compat.cuh | 12 + .../include/common/dcu/kernel_helper.cuh | 342 +++++ infini_train/include/common/dcu/rccl_compat.h | 79 ++ infini_train/include/device.h | 11 + infini_train/include/dtype_dispatch.h | 8 +- infini_train/src/core/ccl/ccl.cc | 7 +- infini_train/src/core/ccl/dcu/rccl_common.cc | 35 + infini_train/src/core/ccl/dcu/rccl_common.h | 37 + infini_train/src/core/ccl/dcu/rccl_impl.cc | 159 +++ infini_train/src/core/ccl/dcu/rccl_impl.h | 51 + .../src/core/runtime/cpu/cpu_guard_impl.cc | 6 +- .../src/core/runtime/dcu/dcu_dispatch.h | 51 + .../src/core/runtime/dcu/dcu_guard_impl.cc | 303 ++++ .../src/core/runtime/dcu/dcu_guard_impl.h | 85 ++ .../core/runtime/dcu/dcu_runtime_common.cc | 58 + .../src/core/runtime/dcu/dcu_runtime_common.h | 64 + infini_train/src/core/runtime/device_guard.cc | 11 +- infini_train/src/device.cc | 28 +- .../src/kernels/dcu/accumulate_grad.hip | 95 ++ infini_train/src/kernels/dcu/cast.hip | 58 + infini_train/src/kernels/dcu/comm.hip | 81 ++ infini_train/src/kernels/dcu/concat.hip | 247 ++++ .../src/kernels/dcu/cross_entropy.hip | 229 +++ infini_train/src/kernels/dcu/elementwise.hip | 1264 +++++++++++++++++ infini_train/src/kernels/dcu/embedding.hip | 126 ++ infini_train/src/kernels/dcu/fill.hip | 47 + infini_train/src/kernels/dcu/gather.hip | 232 +++ infini_train/src/kernels/dcu/layernorm.hip | 208 +++ infini_train/src/kernels/dcu/linear.hip | 504 +++++++ infini_train/src/kernels/dcu/no_op.hip | 30 + infini_train/src/kernels/dcu/outer.hip | 166 +++ infini_train/src/kernels/dcu/reduction.hip | 244 ++++ infini_train/src/kernels/dcu/slice.hip | 209 +++ infini_train/src/kernels/dcu/softmax.hip | 222 +++ infini_train/src/kernels/dcu/split.hip | 182 +++ infini_train/src/kernels/dcu/stack.hip | 160 +++ infini_train/src/kernels/dcu/transform.hip | 593 ++++++++ .../dcu/vocab_parallel_cross_entropy.hip | 127 ++ infini_train/src/nn/parallel/data_parallel.cc | 10 +- infini_train/src/nn/parallel/global.cc | 16 +- .../src/nn/parallel/pp/pipeline_schedule.cc | 16 +- infini_train/src/profiler.cc | 10 +- infini_train/src/tensor.cc | 1 + 51 files changed, 6825 insertions(+), 45 deletions(-) create mode 100644 format create mode 100644 infini_train/include/common/dcu/common_dcu.h create mode 100644 infini_train/include/common/dcu/cub_compat.cuh create mode 100644 infini_train/include/common/dcu/kernel_helper.cuh create mode 100644 infini_train/include/common/dcu/rccl_compat.h create mode 100644 infini_train/src/core/ccl/dcu/rccl_common.cc create mode 100644 infini_train/src/core/ccl/dcu/rccl_common.h create mode 100644 infini_train/src/core/ccl/dcu/rccl_impl.cc create mode 100644 infini_train/src/core/ccl/dcu/rccl_impl.h create mode 100644 infini_train/src/core/runtime/dcu/dcu_dispatch.h create mode 100644 infini_train/src/core/runtime/dcu/dcu_guard_impl.cc create mode 100644 infini_train/src/core/runtime/dcu/dcu_guard_impl.h create mode 100644 infini_train/src/core/runtime/dcu/dcu_runtime_common.cc create mode 100644 infini_train/src/core/runtime/dcu/dcu_runtime_common.h create mode 100644 infini_train/src/kernels/dcu/accumulate_grad.hip create mode 100644 infini_train/src/kernels/dcu/cast.hip create mode 100644 infini_train/src/kernels/dcu/comm.hip create mode 100644 infini_train/src/kernels/dcu/concat.hip create mode 100644 infini_train/src/kernels/dcu/cross_entropy.hip create mode 100644 infini_train/src/kernels/dcu/elementwise.hip create mode 100644 infini_train/src/kernels/dcu/embedding.hip create mode 100644 infini_train/src/kernels/dcu/fill.hip create mode 100644 infini_train/src/kernels/dcu/gather.hip create mode 100644 infini_train/src/kernels/dcu/layernorm.hip create mode 100644 infini_train/src/kernels/dcu/linear.hip create mode 100644 infini_train/src/kernels/dcu/no_op.hip create mode 100644 infini_train/src/kernels/dcu/outer.hip create mode 100644 infini_train/src/kernels/dcu/reduction.hip create mode 100644 infini_train/src/kernels/dcu/slice.hip create mode 100644 infini_train/src/kernels/dcu/softmax.hip create mode 100644 infini_train/src/kernels/dcu/split.hip create mode 100644 infini_train/src/kernels/dcu/stack.hip create mode 100644 infini_train/src/kernels/dcu/transform.hip create mode 100644 infini_train/src/kernels/dcu/vocab_parallel_cross_entropy.hip diff --git a/CMakeLists.txt b/CMakeLists.txt index 6bd8069d..21098ad6 100644 --- a/CMakeLists.txt +++ b/CMakeLists.txt @@ -1,10 +1,25 @@ +<<<<<<< ours cmake_minimum_required(VERSION 3.28) +======= +# Platforms +option(USE_CUDA "Support NVIDIA CUDA" OFF) +option(USE_MACA "Support MetaX MACA" OFF) +option(USE_DCU "Support Hygon DCU through DTK/HIP" OFF) +>>>>>>> theirs option(USE_CUDA "Support NVIDIA CUDA" OFF) option(PROFILE_MODE "ENABLE PROFILE MODE" OFF) option(USE_OMP "Use OpenMP as backend for Eigen" ON) +<<<<<<< ours option(USE_NCCL "Build project for distributed running" ON) option(BUILD_TEST "Build InfiniTrain tests" OFF) +======= +option(USE_NCCL "Build project for distributed running on CUDA using NCCL" ON) +option(USE_MCCL "Build project for distributed running on MACA using MCCL" ON) +option(USE_RCCL "Build project for distributed running on DCU using RCCL" OFF) +option(USE_MPI "Enable MPI for inter-node CPU communication" ON) +cmake_minimum_required(VERSION 3.28) +>>>>>>> theirs project(infini_train VERSION 0.6.0 LANGUAGES CXX) @@ -71,6 +86,11 @@ endif() if(NOT USE_NCCL) list(FILTER SRC EXCLUDE REGEX ".*infini_train/src/core/ccl/cuda/.*") endif() +if(NOT USE_DCU) + list(FILTER SRC EXCLUDE REGEX ".*/(ccl|runtime)/dcu/.*") +elseif(NOT USE_RCCL) + list(FILTER SRC EXCLUDE REGEX ".*/ccl/dcu/.*") +endif() # CPU kernels (*.cc) file(GLOB_RECURSE CPU_KERNELS ${PROJECT_SOURCE_DIR}/infini_train/src/kernels/cpu/*.cc) @@ -127,6 +147,7 @@ endif() # Main framework library # ------------------------------------------------------------------------------ +<<<<<<< ours add_library(infini_train STATIC ${SRC}) target_link_libraries(infini_train PUBLIC @@ -151,6 +172,128 @@ if(USE_CUDA) # keep this. Otherwise it's harmless. target_link_libraries(infini_train PUBLIC nccl) endif() +======= +# ========================= +# DCU backend (Hygon DTK/HIP) +# ========================= +elseif(USE_DCU) + add_compile_definitions(USE_DCU=1) + + set(DCU_PATH "$ENV{DTK_PATH}" CACHE PATH "Hygon DTK installation root") + set(DCU_ARCH "" CACHE STRING "Optional HIP offload architecture reported by rocminfo") + if(NOT DCU_PATH) + set(DCU_PATH /opt/dtk) + endif() + + find_program(HIPCC_EXECUTABLE hipcc + HINTS "${DCU_PATH}/bin" "${DCU_PATH}/llvm/bin" /opt/rocm/bin + REQUIRED) + if(NOT CMAKE_CXX_COMPILER MATCHES "hipcc") + message(WARNING + "DCU kernels must be compiled by hipcc. Reconfigure with " + "-DCMAKE_CXX_COMPILER=${HIPCC_EXECUTABLE}") + endif() + + include_directories("${DCU_PATH}/include") + link_directories("${DCU_PATH}/lib" "${DCU_PATH}/lib64") + + find_library(DCU_RUNTIME_LIB NAMES amdhip64 hip_hcc + HINTS "${DCU_PATH}/lib" "${DCU_PATH}/lib64" /opt/rocm/lib /opt/rocm/lib64 + REQUIRED) + find_library(DCU_BLAS_LIB NAMES hipblas + HINTS "${DCU_PATH}/lib" "${DCU_PATH}/lib64" /opt/rocm/lib /opt/rocm/lib64 + REQUIRED) + + file(GLOB_RECURSE DCU_KERNELS ${PROJECT_SOURCE_DIR}/infini_train/src/kernels/dcu/*.hip) + set_source_files_properties(${DCU_KERNELS} PROPERTIES + LANGUAGE CXX + COMPILE_OPTIONS "-x;hip" + ) + add_library(infini_train_dcu_kernels STATIC ${DCU_KERNELS}) + if(DCU_ARCH) + target_compile_options(infini_train_dcu_kernels PRIVATE "--offload-arch=${DCU_ARCH}") + endif() + target_link_libraries(infini_train_dcu_kernels glog ${DCU_RUNTIME_LIB} ${DCU_BLAS_LIB}) + + add_library(infini_train STATIC ${SRC}) + target_link_libraries(infini_train glog gflags infini_train_cpu_kernels infini_train_dcu_kernels) + + if(USE_RCCL) + message(STATUS "Add USE_RCCL under DCU backend") + find_library(DCU_COMM_LIB NAMES rccl nccl + HINTS + "${DCU_PATH}/lib" + "${DCU_PATH}/lib64" + "${DCU_PATH}/rccl/lib" + "${DCU_PATH}/cuda/cuda-12/targets/x86_64-linux/lib" + "${DCU_PATH}/cuda/targets/x86_64-linux/lib" + /opt/rocm/lib + /opt/rocm/lib64 + REQUIRED) + target_compile_definitions(infini_train PRIVATE USE_RCCL=1) + target_link_libraries(infini_train ${DCU_COMM_LIB}) + endif() + + if(USE_MPI) + target_link_libraries(infini_train ${MPI_LIBS}) + endif() + +# ========================= +# MACA backend (MetaX) +# ========================= +elseif(USE_MACA) + add_compile_definitions(USE_MACA=1) + + # ---- configure MACA SDK paths ---- + # Typical: /opt/maca (can be overridden by -DMACA_PATH=...) + set(MACA_PATH $ENV{MACA_PATH}) + set(CMAKE_C_COMPILER ${MACA_PATH}/mxgpu_llvm/bin/mxcc) + set(CMAKE_CXX_COMPILER ${MACA_PATH}/mxgpu_llvm/bin/mxcc) + + include_directories("${MACA_PATH}/include") + link_directories("${MACA_PATH}/lib") + + # Libraries: mcruntime / mcdnn / mcblas + find_library(MACA_RUNTIME_LIB NAMES mcruntime HINTS "${MACA_PATH}/lib" REQUIRED) + find_library(MACA_DNN_LIB NAMES mcdnn HINTS "${MACA_PATH}/lib" REQUIRED) + find_library(MACA_BLAS_LIB NAMES mcblas HINTS "${MACA_PATH}/lib" REQUIRED) + + file(GLOB_RECURSE MACA_KERNELS ${PROJECT_SOURCE_DIR}/infini_train/src/kernels/maca/*.maca) + set_source_files_properties(${MACA_KERNELS} PROPERTIES + LANGUAGE CXX + COMPILE_OPTIONS "-x;maca" + ) + add_library(infini_train_maca_kernels STATIC ${MACA_KERNELS}) + target_link_libraries(infini_train_maca_kernels glog ${MACA_RUNTIME_LIB} ${MACA_DNN_LIB} ${MACA_BLAS_LIB}) + + add_library(infini_train STATIC ${SRC}) + target_link_libraries(infini_train glog gflags infini_train_cpu_kernels infini_train_maca_kernels) + + if (USE_MCCL) + message(STATUS "Add USE_MCCL under MACA backend, use MCCL (mccl)") + find_library(MACA_COMM_LIB NAMES mccl HINTS "${MACA_PATH}/lib" REQUIRED) + add_compile_definitions(USE_MCCL=1) + target_link_libraries(infini_train ${MACA_COMM_LIB}) + endif() + + if (USE_MPI) + target_link_libraries(infini_train ${MPI_LIBS} Threads::Threads) + + # 有些 MPI 还需要额外 link flags（比如 -Wl,...），也一并带上 + if (MPI_CXX_LINK_FLAGS) + set_target_properties(infini_train PROPERTIES + LINK_FLAGS "${MPI_CXX_LINK_FLAGS}" + ) + endif() + endif() + +# ========================= +# CPU-only backend +# ========================= +else() + add_library(infini_train STATIC ${SRC}) + target_link_libraries(infini_train glog gflags infini_train_cpu_kernels) +>>>>>>> theirs endif() # ------------------------------------------------------------------------------ @@ -168,6 +311,29 @@ function(link_infini_train_exe target_name) "-Wl,--no-whole-archive" "-Wl,--end-group" ) +<<<<<<< ours +======= + elseif(USE_MACA) + target_link_libraries(${target_name} PRIVATE + "-Wl,--start-group" + "-Wl,--whole-archive" + infini_train + infini_train_cpu_kernels + infini_train_maca_kernels + "-Wl,--no-whole-archive" + "-Wl,--end-group" + ) + elseif(USE_DCU) + target_link_libraries(${target_name} PRIVATE + "-Wl,--start-group" + "-Wl,--whole-archive" + infini_train + infini_train_cpu_kernels + infini_train_dcu_kernels + "-Wl,--no-whole-archive" + "-Wl,--end-group" + ) +>>>>>>> theirs else() target_link_libraries(${target_name} PRIVATE "-Wl,--start-group" diff --git a/README.md b/README.md index 1a51d1a8..f9b70224 100644 --- a/README.md +++ b/README.md @@ -35,8 +35,11 @@ A from-scratch C++ training framework for large-scale models with multi-dimensio mkdir build cd build cmake .. -DUSE_CUDA=ON -DUSE_NCCL=ON -make -j -``` +make -j +``` + +For Hygon BW1000 / DCU builds and validation, see +[`docs/dcu_bw1000.md`](docs/dcu_bw1000.md). Build Options: @@ -180,6 +183,7 @@ Multiple parallelism strategies (DDP, TP, SP, PP) can be freely combined to scal Added Autocast, multi-dimensional distributed parallelism (DDP, TP, SP, PP with GPipe / 1F1B / vPP), multi-node training, `no_grad` mode, +<<<<<<< ours and communication–computation overlap with bucketed gradient synchronization. - **2026/06/08** — InfiniTrain **v0.6.0** @@ -208,4 +212,7 @@ Multiple parallelism strategies (DDP, TP, SP, PP) can be freely combined to scal optimizations. Integrated a CTest + GTest based testing infrastructure to strengthen the - framework's automated test workflow. \ No newline at end of file + framework's automated test workflow. +======= + and communication–computation overlap with bucketed gradient synchronization. +>>>>>>> theirs diff --git a/example/gpt2/checkpoint_loader.cc b/example/gpt2/checkpoint_loader.cc index 95e54730..608abe03 100644 --- a/example/gpt2/checkpoint_loader.cc +++ b/example/gpt2/checkpoint_loader.cc @@ -3,6 +3,7 @@ #include #include #include +#include #include #include #include diff --git a/example/gpt2/main.cc b/example/gpt2/main.cc index 60c0c908..cce76eb2 100644 --- a/example/gpt2/main.cc +++ b/example/gpt2/main.cc @@ -75,7 +75,7 @@ DEFINE_uint32(sample_every, 0, "how often to sample from the model?"); // debugging DEFINE_bool(overfit_single_batch, true, "overfit just one batch of data"); // memory management -DEFINE_string(device, "cuda", "device type (cpu/cuda), useless if using parallel training mode"); +DEFINE_string(device, "cuda", "device type (cpu/cuda/maca/dcu), useless if using parallel training mode"); // parallel DEFINE_int32( nthread_per_process, 1, @@ -114,6 +114,11 @@ const std::unordered_set kSupportedModels = {"gpt2", "gpt2-medium", "gpt2-large", "gpt2-xl", "d12", "d24", "d36", "d48"}; constexpr char kDeviceCPU[] = "cpu"; constexpr char kDeviceCUDA[] = "cuda"; +<<<<<<< ours +======= +constexpr char kDeviceMACA[] = "maca"; +constexpr char kDeviceDCU[] = "dcu"; +>>>>>>> theirs constexpr char kDtypeFP32[] = "float32"; constexpr char kDtypeBF16[] = "bfloat16"; const std::unordered_set kSupportedLRDecayStyles @@ -130,11 +135,17 @@ const std::unordered_map kModelToConfigs = { } // namespace DEFINE_validator(model, [](const char *, const std::string &value) { return kSupportedModels.contains(value); }); +<<<<<<< ours DEFINE_validator(device, [](const char *, const std::string &value) { return value == kDeviceCPU || value == kDeviceCUDA; }); DEFINE_validator(zero_stage, [](const char *, int32_t value) { return value >= 0 && value <= 3; }); DEFINE_validator(lr_decay_style, [](const char *, const std::string &value) { return kSupportedLRDecayStyles.contains(value); }); +======= +DEFINE_validator(device, [](const char *, const std::string &value) { + return value == kDeviceCPU || value == kDeviceCUDA || value == kDeviceMACA || value == kDeviceDCU; +}); +>>>>>>> theirs void Train(const nn::parallel::Rank &rank) { using namespace nn::parallel; @@ -181,7 +192,18 @@ void Train(const nn::parallel::Rank &rank) { const ProcessGroup *pp_pg = nullptr; if (rank.IsParallel()) { +<<<<<<< ours device = Device(Device::DeviceType::kCUDA, rank.thread_rank()); +======= + auto parallel_device_type = Device::DeviceType::kCUDA; + if (FLAGS_device == kDeviceMACA) { + parallel_device_type = Device::DeviceType::kMACA; + } else if (FLAGS_device == kDeviceDCU) { + parallel_device_type = Device::DeviceType::kDCU; + } + device = Device(parallel_device_type, rank.thread_rank()); + +>>>>>>> theirs auto *pg_factory = ProcessGroupFactory::Instance(device.type()); if (ddp_world_size > 1) { @@ -206,7 +228,19 @@ void Train(const nn::parallel::Rank &rank) { nn::parallel::pp_rank = pp_rank; } } else { +<<<<<<< ours device = FLAGS_device == kDeviceCPU ? Device() : Device(Device::DeviceType::kCUDA, 0); +======= + if (FLAGS_device == kDeviceCPU) { + device = Device(); + } else if (FLAGS_device == kDeviceMACA) { + device = Device(Device::DeviceType::kMACA, 0); + } else if (FLAGS_device == kDeviceDCU) { + device = Device(Device::DeviceType::kDCU, 0); + } else { + device = Device(Device::DeviceType::kCUDA, 0); + } +>>>>>>> theirs } // calculate gradient accumulation from the desired total batch size and the current run configuration diff --git a/example/llama3/checkpoint_loader.cc b/example/llama3/checkpoint_loader.cc index f3590af6..e8425925 100644 --- a/example/llama3/checkpoint_loader.cc +++ b/example/llama3/checkpoint_loader.cc @@ -3,6 +3,7 @@ #include #include #include +#include #include #include #include diff --git a/example/llama3/main.cc b/example/llama3/main.cc index 302e0808..e41703cd 100644 --- a/example/llama3/main.cc +++ b/example/llama3/main.cc @@ -74,7 +74,7 @@ DEFINE_uint32(sample_every, 0, "how often to sample from the model?"); // debugging DEFINE_bool(overfit_single_batch, true, "overfit just one batch of data"); // memory management -DEFINE_string(device, "cuda", "device type (cpu/cuda), useless if using parallel training mode"); +DEFINE_string(device, "cuda", "device type (cpu/cuda/maca/dcu), useless if using parallel training mode"); // parallel DEFINE_int32( nthread_per_process, 1, @@ -110,6 +110,11 @@ namespace { const std::unordered_set kSupportedModels = {"llama3"}; constexpr char kDeviceCPU[] = "cpu"; constexpr char kDeviceCUDA[] = "cuda"; +<<<<<<< ours +======= +constexpr char kDeviceMACA[] = "maca"; +constexpr char kDeviceDCU[] = "dcu"; +>>>>>>> theirs constexpr char kDtypeFP32[] = "float32"; constexpr char kDtypeBF16[] = "bfloat16"; const std::unordered_set kSupportedLRDecayStyles @@ -117,11 +122,17 @@ const std::unordered_set kSupportedLRDecayStyles } // namespace DEFINE_validator(model, [](const char *, const std::string &value) { return kSupportedModels.contains(value); }); +<<<<<<< ours DEFINE_validator(device, [](const char *, const std::string &value) { return value == kDeviceCPU || value == kDeviceCUDA; }); DEFINE_validator(zero_stage, [](const char *, int32_t value) { return value >= 0 && value <= 3; }); DEFINE_validator(lr_decay_style, [](const char *, const std::string &value) { return kSupportedLRDecayStyles.contains(value); }); +======= +DEFINE_validator(device, [](const char *, const std::string &value) { + return value == kDeviceCPU || value == kDeviceCUDA || value == kDeviceMACA || value == kDeviceDCU; +}); +>>>>>>> theirs void Train(const nn::parallel::Rank &rank) { using namespace nn::parallel; @@ -167,7 +178,18 @@ void Train(const nn::parallel::Rank &rank) { const ProcessGroup *pp_pg = nullptr; if (rank.IsParallel()) { +<<<<<<< ours device = Device(Device::DeviceType::kCUDA, rank.thread_rank()); +======= + auto parallel_device_type = Device::DeviceType::kCUDA; + if (FLAGS_device == kDeviceMACA) { + parallel_device_type = Device::DeviceType::kMACA; + } else if (FLAGS_device == kDeviceDCU) { + parallel_device_type = Device::DeviceType::kDCU; + } + device = Device(parallel_device_type, rank.thread_rank()); + +>>>>>>> theirs auto *pg_factory = ProcessGroupFactory::Instance(device.type()); if (ddp_world_size > 1) { @@ -192,7 +214,19 @@ void Train(const nn::parallel::Rank &rank) { nn::parallel::pp_rank = pp_rank; } } else { +<<<<<<< ours device = FLAGS_device == kDeviceCPU ? Device() : Device(Device::DeviceType::kCUDA, 0); +======= + if (FLAGS_device == kDeviceCPU) { + device = Device(); + } else if (FLAGS_device == kDeviceMACA) { + device = Device(Device::DeviceType::kMACA, 0); + } else if (FLAGS_device == kDeviceDCU) { + device = Device(Device::DeviceType::kDCU, 0); + } else { + device = Device(Device::DeviceType::kCUDA, 0); + } +>>>>>>> theirs } // calculate gradient accumulation from the desired total batch size and the current run configuration diff --git a/format b/format new file mode 100644 index 00000000..ede9c297 --- /dev/null +++ b/format @@ -0,0 +1,147 @@ +#pragma once + +#include +#include +#include +#include +#include +#include +#include + +namespace std { +namespace infini_train_format_compat { + +struct FormatSpec { + int width = 0; + int precision = -1; + bool left = false; + char type = '\0'; +}; + +inline FormatSpec ParseSpec(std::string_view spec) { + FormatSpec parsed; + if (!spec.empty() && spec.front() == ':') { + spec.remove_prefix(1); + } + if (!spec.empty() && spec.front() == '<') { + parsed.left = true; + spec.remove_prefix(1); + } else if (!spec.empty() && spec.front() == '>') { + spec.remove_prefix(1); + } + + while (!spec.empty() && spec.front() >= '0' && spec.front() <= '9') { + parsed.width = parsed.width * 10 + (spec.front() - '0'); + spec.remove_prefix(1); + } + + if (!spec.empty() && spec.front() == '.') { + spec.remove_prefix(1); + parsed.precision = 0; + while (!spec.empty() && spec.front() >= '0' && spec.front() <= '9') { + parsed.precision = parsed.precision * 10 + (spec.front() - '0'); + spec.remove_prefix(1); + } + } + + if (!spec.empty()) { + parsed.type = spec.front(); + } + return parsed; +} + +template std::string FormatOne(std::string_view spec_text, T &&value) { + const auto spec = ParseSpec(spec_text); + std::ostringstream oss; + if (spec.left) { + oss << std::left; + } + if (spec.width > 0) { + oss << std::setw(spec.width); + } + if (spec.precision >= 0) { + oss << std::setprecision(spec.precision); + } + if (spec.type == 'f') { + oss << std::fixed; + } else if (spec.type == 'e') { + oss << std::scientific; + } + oss << std::forward(value); + return oss.str(); +} + +inline void CollectFormatArgs(std::vector &) {} + +template +void CollectFormatArgs(std::vector &out, std::string_view spec, T &&value, Rest &&...rest) { + out.push_back(FormatOne(spec, std::forward(value))); + if constexpr (sizeof...(Rest) > 0) { + CollectFormatArgs(out, "", std::forward(rest)...); + } +} + +template void PushFormattedArg(std::vector &out, std::string_view spec, T &&value) { + out.push_back(FormatOne(spec, std::forward(value))); +} + +template std::vector MakeFormatArgs(std::string_view fmt, Args &&...args) { + std::vector result; + result.reserve(sizeof...(Args)); + + size_t arg_index = 0; + auto specs = std::vector{}; + for (size_t i = 0; i < fmt.size(); ++i) { + if (fmt[i] != '{' || (i + 1 < fmt.size() && fmt[i + 1] == '{')) { + if (fmt[i] == '{') { + ++i; + } + continue; + } + const size_t close = fmt.find('}', i + 1); + if (close == std::string_view::npos) { + break; + } + specs.push_back(fmt.substr(i + 1, close - i - 1)); + i = close; + ++arg_index; + } + + size_t spec_index = 0; + (PushFormattedArg(result, spec_index < specs.size() ? specs[spec_index++] : std::string_view{}, std::forward(args)), ...); + return result; +} + +} // namespace infini_train_format_compat + +template std::string format(std::string_view fmt, Args &&...args) { + const auto formatted_args = infini_train_format_compat::MakeFormatArgs(fmt, std::forward(args)...); + std::ostringstream out; + size_t arg_index = 0; + + for (size_t i = 0; i < fmt.size(); ++i) { + if (fmt[i] == '{') { + if (i + 1 < fmt.size() && fmt[i + 1] == '{') { + out << '{'; + ++i; + continue; + } + const size_t close = fmt.find('}', i + 1); + if (close != std::string_view::npos) { + if (arg_index < formatted_args.size()) { + out << formatted_args[arg_index++]; + } + i = close; + continue; + } + } else if (fmt[i] == '}' && i + 1 < fmt.size() && fmt[i + 1] == '}') { + out << '}'; + ++i; + continue; + } + out << fmt[i]; + } + return out.str(); +} + +} // namespace std diff --git a/infini_train/include/common/dcu/common_dcu.h b/infini_train/include/common/dcu/common_dcu.h new file mode 100644 index 00000000..6c200f2f --- /dev/null +++ b/infini_train/include/common/dcu/common_dcu.h @@ -0,0 +1,46 @@ +#pragma once + +#include +#include +#include +#ifdef USE_RCCL +#include "infini_train/include/common/dcu/rccl_compat.h" +#endif + +#include "glog/logging.h" + +#include "infini_train/include/common/dcu/cub_compat.cuh" +#include "infini_train/include/common/dcu/kernel_helper.cuh" + +namespace infini_train::common::dcu { + +// Common HIP Macros +#define HIP_CHECK(call) \ + do { \ + hipError_t status = call; \ + if (status != hipSuccess) { \ + LOG(FATAL) << "HIP Error: " << hipGetErrorString(status) << " at " << __FILE__ << ":" << __LINE__; \ + } \ + } while (0) + +#define HIPBLAS_CHECK(call) \ + do { \ + hipblasStatus_t status = call; \ + if (status != HIPBLAS_STATUS_SUCCESS) { \ + LOG(FATAL) << "HIPBLAS Error: status=" << static_cast(status) << " at " << __FILE__ << ":" \ + << __LINE__; \ + } \ + } while (0) + +#ifdef USE_RCCL +#define NCCL_CHECK(expr) \ + do { \ + ncclResult_t _status = (expr); \ + if (_status != ncclSuccess) { \ + LOG(FATAL) << "NCCL error: " << ncclGetErrorString(_status) << " at " << __FILE__ << ":" << __LINE__ \ + << " (" << #expr << ")"; \ + } \ + } while (0) +#endif + +} // namespace infini_train::common::dcu diff --git a/infini_train/include/common/dcu/cub_compat.cuh b/infini_train/include/common/dcu/cub_compat.cuh new file mode 100644 index 00000000..fd1ab9ae --- /dev/null +++ b/infini_train/include/common/dcu/cub_compat.cuh @@ -0,0 +1,12 @@ +#pragma once + +#include +#include + +namespace infini_train::kernels::dcu { + +using CubSumOp = hipcub::Sum; +using CubMaxOp = hipcub::Max; +using CubMinOp = hipcub::Min; + +} // namespace infini_train::kernels::dcu diff --git a/infini_train/include/common/dcu/kernel_helper.cuh b/infini_train/include/common/dcu/kernel_helper.cuh new file mode 100644 index 00000000..262fd253 --- /dev/null +++ b/infini_train/include/common/dcu/kernel_helper.cuh @@ -0,0 +1,342 @@ +#pragma once +/* +#include +#include +#include +#include +*/ +#include +#include + +namespace infini_train::common::dcu { + +template __device__ __forceinline__ void AtomicAdd(T *address, T value) { atomicAdd(address, value); } + +// HIP lacks a portable BF16 atomicAdd on older DCU targets. Update the +// containing 32-bit word with CAS so adjacent BF16 values remain intact. +__device__ __forceinline__ void AtomicAdd(hip_bfloat16 *address, hip_bfloat16 value) { + const auto raw_address = reinterpret_cast(address); + auto *base = reinterpret_cast(raw_address & ~std::uintptr_t{0x3}); + const bool upper = (raw_address & 0x2) != 0; + + unsigned int old = *base; + unsigned int assumed = 0; + do { + assumed = old; + hip_bfloat16 current; + current.data = static_cast(upper ? (assumed >> 16) : (assumed & 0xffff)); + hip_bfloat16 updated(static_cast(current) + static_cast(value)); + const unsigned int updated_bits = static_cast(updated.data); + const unsigned int replacement + = upper ? ((assumed & 0x0000ffffU) | (updated_bits << 16)) + : ((assumed & 0xffff0000U) | updated_bits); + old = atomicCAS(base, assumed, replacement); + } while (old != assumed); +} + +/** + * Converts a value between arbitrary types with specialized handling for + * HIP floating-point precisions. For primitive types, this offers perfect + * forwarding which preserves value categories (lvalues/rvalues) + * + * @tparam DST Destination type (deduced) + * @tparam SRC Source type (deduced) + * @param x Input value (preserves const/volatile and value category) + * @return Value converted to DST type + * + * Example: + * half h = Cast(3.14f); // float -> half (HIP intrinsic) + * float f = Cast(h); // half -> float (HIP intrinsic) + * int i = Cast(2.718); // double -> int (standard cast) + */ +// TODO(lzm): add support for half and hip_bfloat16 conversions with integral types +template __host__ __device__ DST Cast(SRC &&x) { + static_assert(!std::is_reference_v, "Cast cannot return reference types"); + + using SRC_base = std::remove_cv_t>; + using DST_base = std::remove_cv_t>; + + // hip_bfloat16 conversions + if constexpr (std::is_same_v) { + if constexpr (std::is_same_v) { + return static_cast(x); + } else if constexpr (std::is_same_v) { + return static_cast(static_cast(x)); + } else if constexpr (std::is_same_v) { + return __float2half(static_cast(x)); + } + } + // half conversions + else if constexpr (std::is_same_v) { + if constexpr (std::is_same_v) { + return __half2float(x); + } else if constexpr (std::is_same_v) { + return static_cast(__half2float(x)); + } else if constexpr (std::is_same_v) { + return hip_bfloat16(__half2float(x)); + } + } + // float conversions to reduced precision + else if constexpr (std::is_same_v) { + if constexpr (std::is_same_v) { + return hip_bfloat16(x); + } else if constexpr (std::is_same_v) { + return __float2half(x); + } + } + // double conversions to reduced precision + else if constexpr (std::is_same_v) { + if constexpr (std::is_same_v) { + return hip_bfloat16(static_cast(x)); + } else if constexpr (std::is_same_v) { + return __double2half(x); + } + } + // Fallback for all other conversions + return (DST)(std::forward(x)); +} + +template __device__ __forceinline__ T Neg(const T &x) { + if constexpr (std::is_same_v) { + return hip_bfloat16(-static_cast(x)); + } else if constexpr (std::is_same_v) { + return __hneg(x); + } else { + return -x; + } +} + +template __device__ __forceinline__ T Reciprocal(const T &x) { + if constexpr (std::is_same_v) { + return __hdiv(__float2half(1.0f), x); + } else if constexpr (std::is_same_v) { + return hip_bfloat16(1.0f / static_cast(x)); + } else { + return T(1) / x; + } +} + +template __device__ __forceinline__ T Sin(const T &x) { + if constexpr (std::is_same_v) { + return __float2half(__sinf(__half2float(x))); + } else if constexpr (std::is_same_v) { + return hip_bfloat16(__sinf(static_cast(x))); + } else if constexpr (std::is_same_v) { + return __sinf(x); + } else { + return std::sin(x); + } +} + +template __device__ __forceinline__ T Cos(const T &x) { + if constexpr (std::is_same_v) { + return __float2half(__cosf(__half2float(x))); + } else if constexpr (std::is_same_v) { + return hip_bfloat16(__cosf(static_cast(x))); + } else if constexpr (std::is_same_v) { + return __cosf(x); + } else { + return std::cos(x); + } +} + +template __device__ __forceinline__ T Tanh(const T &x) { + if constexpr (std::is_same_v) { + return hip_bfloat16(tanhf(static_cast(x))); + } else if constexpr (std::is_same_v) { + return htanh(x); + } else if constexpr (std::is_same_v) { + return tanhf(x); + } else { + return std::tanh(x); + } +} + +template __device__ __forceinline__ T Pow(const T &x, const T &exponent) { + if constexpr (std::is_same_v) { + float x_ = static_cast(x); + float exponent_ = static_cast(exponent); + float ans_f = __powf(x_, exponent_); + return hip_bfloat16(__isnan(ans_f) ? std::pow(x_, exponent_) : ans_f); + } else if constexpr (std::is_same_v) { + float x_ = __half2float(x); + float exponent_ = __half2float(exponent); + float ans_f = __powf(x_, exponent_); + return __float2half(__isnan(ans_f) ? std::pow(x_, exponent_) : ans_f); + } else if constexpr (std::is_same_v) { + return powf(x, exponent); + } else { + return std::pow(x, exponent); + } +} + +template __device__ __forceinline__ T Rsqrt(const T &x) { + if constexpr (std::is_same_v) { + return __float2half(rsqrtf(__half2float(x))); + } else if constexpr (std::is_same_v) { + return hip_bfloat16(rsqrtf(static_cast(x))); + } else if constexpr (std::is_same_v) { + return rsqrtf(x); + } else { + return T(1) / std::sqrt(T(x)); + } +} + +template __device__ __forceinline__ T Exp(const T &x) { + if constexpr (std::is_same_v) { + return hip_bfloat16(__expf(static_cast(x))); + } else if constexpr (std::is_same_v) { + return hexp(x); + } else if constexpr (std::is_same_v) { + return __expf(x); + } else { + return std::exp(x); + } +} + +template __device__ __forceinline__ T Log(const T &x) { + if constexpr (std::is_same_v) { + return hip_bfloat16(__logf(static_cast(x))); + } else if constexpr (std::is_same_v) { + return __float2half(__logf(__half2float(x))); + } else if constexpr (std::is_same_v) { + return __logf(x); + } else { + return std::log(x); + } +} + +template __device__ __forceinline__ T Add(const T &a, const T &b) { + if constexpr (std::is_same_v) { + return hip_bfloat16(static_cast(a) + static_cast(b)); + } else if constexpr (std::is_same_v) { + return __hadd(a, b); + } else { + return a + b; + } +} + +template __device__ __forceinline__ T Sub(const T &a, const T &b) { + if constexpr (std::is_same_v) { + return hip_bfloat16(static_cast(a) - static_cast(b)); + } else if constexpr (std::is_same_v) { + return __hsub(a, b); + } else { + return a - b; + } +} + +template __device__ __forceinline__ T Mul(const T &a, const T &b) { + if constexpr (std::is_same_v) { + return hip_bfloat16(static_cast(a) * static_cast(b)); + } else if constexpr (std::is_same_v) { + return __hmul(a, b); + } else { + return a * b; + } +} + +template __device__ __forceinline__ T Div(const T &a, const T &b) { + if constexpr (std::is_same_v) { + return hip_bfloat16(static_cast(a) / static_cast(b)); + } else if constexpr (std::is_same_v) { + return __hdiv(a, b); + } else { + return a / b; + } +} + +template __device__ __forceinline__ T Sigmoid(const T &x) { + if constexpr (std::is_same_v) { + return 1.0f / (1.0f + expf(-x)); + } else if constexpr (std::is_same_v) { + const float xf = static_cast(x); + return hip_bfloat16(1.0f / (1.0f + expf(-xf))); + } else if constexpr (std::is_same_v) { + return __hdiv(T(1), T(1) + hexp(-x)); + } else { + return T(1) / (T(1) + std::exp(-x)); + } +} + +template __device__ __forceinline__ T Max(const T &a, const T &b) { + if constexpr (std::is_same_v) { + return static_cast(a) <= static_cast(b) ? b : a; + } else if constexpr (std::is_same_v) { + return __hle(a, b) ? b : a; + } else if constexpr (std::is_same_v) { + return fmaxf(a, b); + } else { + return std::max(a, b); + } +} + +template __device__ __forceinline__ T Min(const T &a, const T &b) { + if constexpr (std::is_same_v) { + return static_cast(a) <= static_cast(b) ? a : b; + } else if constexpr (std::is_same_v) { + return __hle(a, b) ? a : b; + } else if constexpr (std::is_same_v) { + return fminf(a, b); + } else { + return std::min(a, b); + } +} + +template __device__ __forceinline__ T Fma(const T &x, const T &y, const T &z) { + if constexpr (std::is_same_v) { + return __hfma(x, y, z); + } else if constexpr (std::is_same_v) { + return hip_bfloat16(__fmaf_rn(static_cast(x), static_cast(y), static_cast(z))); + } else if constexpr (std::is_same_v) { + return __fmaf_rn(x, y, z); + } else { + return std::fma(x, y, z); + } +} + +template ::value> * = nullptr> +__device__ __forceinline__ void fastSpecializedAtomicAdd(scalar_t *tensor, index_t index, const index_t num_elements, + scalar_t value) { + __half *target_addr = tensor + index; + bool low_byte = ((reinterpret_cast(target_addr) & (sizeof(__half2) - 1)) == 0); + + if (low_byte && index < (num_elements - 1)) { + __half2 value2 = __halves2half2(value, __float2half(0.0f)); + atomicAdd(reinterpret_cast<__half2 *>(target_addr), value2); + + } else if (!low_byte && index > 0) { + __half2 value2 = __halves2half2(__float2half(0.0f), value); + atomicAdd(reinterpret_cast<__half2 *>(target_addr - 1), value2); + + } else { + atomicAdd(target_addr, value); + } +} + +template ::value> * = nullptr> +__device__ __forceinline__ void fastSpecializedAtomicAdd(scalar_t *tensor, index_t index, const index_t num_elements, + scalar_t value) { + AtomicAdd(tensor + index, value); +} + +template ::value + && !std::is_same::value> * = nullptr> +__device__ __forceinline__ void fastSpecializedAtomicAdd(scalar_t *tensor, index_t index, + const index_t /*num_elements*/, scalar_t value) { + atomicAdd(tensor + index, value); +} + +template +__device__ __forceinline__ void fastAtomicAdd(scalar_t *tensor, index_t index, const index_t num_elements, + scalar_t value, bool fast_atomics) { + if (fast_atomics) { + fastSpecializedAtomicAdd(tensor, index, num_elements, value); + } else { + AtomicAdd(tensor + index, value); + } +} +} // namespace infini_train::common::dcu diff --git a/infini_train/include/common/dcu/rccl_compat.h b/infini_train/include/common/dcu/rccl_compat.h new file mode 100644 index 00000000..b1e418f0 --- /dev/null +++ b/infini_train/include/common/dcu/rccl_compat.h @@ -0,0 +1,79 @@ +#pragma once + +#include +#include + +#include + +#define NCCL_UNIQUE_ID_BYTES 128 + +typedef struct ncclComm *ncclComm_t; + +typedef struct { + char internal[NCCL_UNIQUE_ID_BYTES]; +} ncclUniqueId; + +typedef enum { + ncclSuccess = 0, + ncclUnhandledCudaError = 1, + ncclSystemError = 2, + ncclInternalError = 3, + ncclInvalidArgument = 4, + ncclInvalidUsage = 5, + ncclRemoteError = 6, + ncclInProgress = 7, + ncclNumResults = 8 +} ncclResult_t; + +typedef enum { + ncclInt8 = 0, + ncclChar = 0, + ncclUint8 = 1, + ncclInt32 = 2, + ncclInt = 2, + ncclUint32 = 3, + ncclInt64 = 4, + ncclUint64 = 5, + ncclFloat16 = 6, + ncclHalf = 6, + ncclFloat32 = 7, + ncclFloat = 7, + ncclFloat64 = 8, + ncclDouble = 8, + ncclBfloat16 = 9, + ncclNumTypes = 10 +} ncclDataType_t; + +typedef enum { + ncclSum = 0, + ncclProd = 1, + ncclMax = 2, + ncclMin = 3, + ncclAvg = 4, + ncclNumOps = 5 +} ncclRedOp_t; + +extern "C" { +const char *ncclGetErrorString(ncclResult_t result); +ncclResult_t ncclGetUniqueId(ncclUniqueId *uniqueId); +ncclResult_t ncclCommInitRank(ncclComm_t *comm, int nranks, ncclUniqueId commId, int rank); +ncclResult_t ncclCommInitAll(ncclComm_t *comms, int ndev, const int *devlist); +ncclResult_t ncclCommDestroy(ncclComm_t comm); +ncclResult_t ncclCommGetAsyncError(ncclComm_t comm, ncclResult_t *asyncError); +ncclResult_t ncclGroupStart(); +ncclResult_t ncclGroupEnd(); +ncclResult_t ncclAllReduce(const void *sendbuff, void *recvbuff, size_t count, ncclDataType_t datatype, + ncclRedOp_t op, ncclComm_t comm, hipStream_t stream); +ncclResult_t ncclBroadcast(const void *sendbuff, void *recvbuff, size_t count, ncclDataType_t datatype, int root, + ncclComm_t comm, hipStream_t stream); +ncclResult_t ncclReduce(const void *sendbuff, void *recvbuff, size_t count, ncclDataType_t datatype, ncclRedOp_t op, + int root, ncclComm_t comm, hipStream_t stream); +ncclResult_t ncclAllGather(const void *sendbuff, void *recvbuff, size_t sendcount, ncclDataType_t datatype, + ncclComm_t comm, hipStream_t stream); +ncclResult_t ncclReduceScatter(const void *sendbuff, void *recvbuff, size_t recvcount, ncclDataType_t datatype, + ncclRedOp_t op, ncclComm_t comm, hipStream_t stream); +ncclResult_t ncclSend(const void *sendbuff, size_t count, ncclDataType_t datatype, int peer, ncclComm_t comm, + hipStream_t stream); +ncclResult_t ncclRecv(void *recvbuff, size_t count, ncclDataType_t datatype, int peer, ncclComm_t comm, + hipStream_t stream); +} diff --git a/infini_train/include/device.h b/infini_train/include/device.h index 28db395f..8942d9e5 100644 --- a/infini_train/include/device.h +++ b/infini_train/include/device.h @@ -13,7 +13,13 @@ class Device { enum class DeviceType : int8_t { kCPU = 0, kCUDA = 1, +<<<<<<< ours kCount = 2, +======= + kMACA = 2, + kDCU = 3, + kCount = 4, +>>>>>>> theirs kInvalid = -1, }; @@ -30,6 +36,11 @@ class Device { bool IsCPU() const; bool IsCUDA() const; +<<<<<<< ours +======= + bool IsMACA() const; + bool IsDCU() const; +>>>>>>> theirs std::string ToString() const; diff --git a/infini_train/include/dtype_dispatch.h b/infini_train/include/dtype_dispatch.h index 8bd5054b..1847fcc5 100644 --- a/infini_train/include/dtype_dispatch.h +++ b/infini_train/include/dtype_dispatch.h @@ -1,6 +1,6 @@ #pragma once -#include +#include #include #include #include @@ -325,8 +325,10 @@ auto DispatchByTypeMap(const std::vector &dtypes, Functor &&func, std: constexpr size_t kNumLists = sizeof...(AllowedTypeLists); if (dtypes.size() != kNumLists) { - LOG(FATAL) << std::format("DispatchByTypeMap expects {} dtypes, but only got {} in {}", kNumLists, - dtypes.size(), context_identifier); + std::ostringstream oss; + oss << "DispatchByTypeMap expects " << kNumLists << " dtypes, but only got " << dtypes.size() << " in " + << context_identifier; + LOG(FATAL) << oss.str(); std::abort(); } diff --git a/infini_train/src/core/ccl/ccl.cc b/infini_train/src/core/ccl/ccl.cc index 92c14cc6..743e4cde 100644 --- a/infini_train/src/core/ccl/ccl.cc +++ b/infini_train/src/core/ccl/ccl.cc @@ -1,6 +1,5 @@ #include "infini_train/include/core/ccl/ccl.h" -#include #include #include @@ -74,12 +73,12 @@ CclImplRegistry &CclImplRegistry::Instance() { void CclImplRegistry::Register(Device::DeviceType type, std::unique_ptr impl) { if (type != impl->Type()) { - LOG(FATAL) << std::format("Register CclImpl with type {}, but as type {}", static_cast(impl->Type()), - static_cast(type)); + LOG(FATAL) << "Register CclImpl with type " << static_cast(impl->Type()) << ", but as type " + << static_cast(type); } if (impls_.contains(type)) { - LOG(FATAL) << std::format("CclImpl for type {} already registered", static_cast(type)); + LOG(FATAL) << "CclImpl for type " << static_cast(type) << " already registered"; } impls_[type] = std::move(impl); diff --git a/infini_train/src/core/ccl/dcu/rccl_common.cc b/infini_train/src/core/ccl/dcu/rccl_common.cc new file mode 100644 index 00000000..8cbb1cd9 --- /dev/null +++ b/infini_train/src/core/ccl/dcu/rccl_common.cc @@ -0,0 +1,35 @@ +#include "infini_train/src/core/ccl/dcu/rccl_common.h" + +#include + +#include "glog/logging.h" + +namespace infini_train::core { + +RcclComm::RcclComm() = default; + +RcclComm::RcclComm(ncclComm_t comm) : nccl_comm_(comm) {} + +ncclComm_t RcclComm::nccl_comm() const { return nccl_comm_; } + +void RcclComm::set_nccl_comm(ncclComm_t comm) { nccl_comm_ = comm; } + +RcclUniqueId::RcclUniqueId() = default; + +RcclUniqueId::RcclUniqueId(const ncclUniqueId &id) : id_(id) {} + +size_t RcclUniqueId::Size() const { return sizeof(id_); } + +const void *RcclUniqueId::Data() const { return &id_; } + +void RcclUniqueId::Load(const void *src, size_t size) { + CHECK_NOTNULL(src); + CHECK_EQ(size, sizeof(id_)); + std::memcpy(&id_, src, sizeof(id_)); +} + +ncclUniqueId *RcclUniqueId::nccl_unique_id() { return &id_; } + +const ncclUniqueId *RcclUniqueId::nccl_unique_id() const { return &id_; } + +} // namespace infini_train::core diff --git a/infini_train/src/core/ccl/dcu/rccl_common.h b/infini_train/src/core/ccl/dcu/rccl_common.h new file mode 100644 index 00000000..08ef0709 --- /dev/null +++ b/infini_train/src/core/ccl/dcu/rccl_common.h @@ -0,0 +1,37 @@ +#pragma once + +#include "infini_train/include/common/dcu/rccl_compat.h" + +#include "infini_train/include/core/ccl/ccl_common.h" + +namespace infini_train::core { + +class RcclComm final : public CclComm { +public: + RcclComm(); + explicit RcclComm(ncclComm_t comm); + + ncclComm_t nccl_comm() const; + void set_nccl_comm(ncclComm_t comm); + +private: + ncclComm_t nccl_comm_ = nullptr; +}; + +class RcclUniqueId final : public CclUniqueId { +public: + RcclUniqueId(); + explicit RcclUniqueId(const ncclUniqueId &id); + + size_t Size() const override; + const void *Data() const override; + void Load(const void *src, size_t size) override; + + ncclUniqueId *nccl_unique_id(); + const ncclUniqueId *nccl_unique_id() const; + +private: + ncclUniqueId id_; +}; + +} // namespace infini_train::core diff --git a/infini_train/src/core/ccl/dcu/rccl_impl.cc b/infini_train/src/core/ccl/dcu/rccl_impl.cc new file mode 100644 index 00000000..2bfa4228 --- /dev/null +++ b/infini_train/src/core/ccl/dcu/rccl_impl.cc @@ -0,0 +1,159 @@ +#include "infini_train/src/core/ccl/dcu/rccl_impl.h" + +#include + +#include "glog/logging.h" + +#include "infini_train/include/common/dcu/common_dcu.h" +#include "infini_train/include/core/runtime/runtime_common.h" +#include "infini_train/include/device.h" + +#include "infini_train/src/core/ccl/dcu/rccl_common.h" +#include "infini_train/src/core/runtime/dcu/dcu_runtime_common.h" + +namespace infini_train::core::dcu { +namespace { + +inline const std::unordered_map kRcclDtypeMap = { + {DataType::kUINT8, ncclUint8}, {DataType::kINT8, ncclInt8}, {DataType::kUINT32, ncclUint32}, + {DataType::kINT32, ncclInt32}, {DataType::kUINT64, ncclUint64}, {DataType::kINT64, ncclInt64}, + {DataType::kBFLOAT16, ncclBfloat16}, {DataType::kFLOAT16, ncclHalf}, {DataType::kFLOAT32, ncclFloat32}, + {DataType::kFLOAT64, ncclFloat64}, +}; + +inline const std::unordered_map kRcclReduceOpMap = { + {nn::parallel::function::ReduceOpType::kSum, ncclSum}, {nn::parallel::function::ReduceOpType::kProd, ncclProd}, + {nn::parallel::function::ReduceOpType::kMin, ncclMin}, {nn::parallel::function::ReduceOpType::kMax, ncclMax}, + {nn::parallel::function::ReduceOpType::kAvg, ncclAvg}, +}; + +inline ncclComm_t GetRcclComm(const CclComm *comm) { + auto *nccl_comm = dynamic_cast(comm); + CHECK_NOTNULL(nccl_comm); + return nccl_comm->nccl_comm(); +} + +inline void SetRcclComm(CclComm *comm, ncclComm_t nccl_comm) { + auto *typed_comm = dynamic_cast(comm); + CHECK_NOTNULL(typed_comm); + typed_comm->set_nccl_comm(nccl_comm); +} + +inline const ncclUniqueId &GetRcclUniqueId(const CclUniqueId &unique_id) { + auto *nccl_unique_id = dynamic_cast(&unique_id); + CHECK_NOTNULL(nccl_unique_id); + return *nccl_unique_id->nccl_unique_id(); +} + +inline hipStream_t GetDcuStream(Stream *stream) { + auto *cuda_stream = dynamic_cast(stream); + CHECK_NOTNULL(cuda_stream); + return cuda_stream->hip_stream(); +} + +} // namespace + +Device::DeviceType RcclImpl::Type() const { return Device::DeviceType::kDCU; } + +void RcclImpl::GroupStart() const { NCCL_CHECK(ncclGroupStart()); } + +void RcclImpl::GroupEnd() const { NCCL_CHECK(ncclGroupEnd()); } + +void RcclImpl::GetAsyncError(const CclComm *comm, CclStatus *async_error) const { + ncclResult_t nccl_async_error = ncclSuccess; + NCCL_CHECK(ncclCommGetAsyncError(GetRcclComm(comm), &nccl_async_error)); + if (async_error != nullptr) { + *async_error = (nccl_async_error == ncclSuccess) ? CclStatus::kSuccess : CclStatus::kError; + } +} + +void RcclImpl::GetUniqueId(CclUniqueId **unique_id) const { + CHECK_NOTNULL(unique_id); + if (*unique_id == nullptr) { + *unique_id = new RcclUniqueId(); + } + auto *nccl_unique_id = dynamic_cast(*unique_id); + CHECK_NOTNULL(nccl_unique_id); + NCCL_CHECK(ncclGetUniqueId(nccl_unique_id->nccl_unique_id())); +} + +void RcclImpl::CommInitAll(CclComm **comms, int ndev, const int *devlist) const { + CHECK_NOTNULL(comms); + CHECK_GT(ndev, 0); + CHECK_NOTNULL(devlist); + + std::vector nccl_comms(static_cast(ndev), nullptr); + NCCL_CHECK(ncclCommInitAll(nccl_comms.data(), ndev, devlist)); + for (int i = 0; i < ndev; ++i) { + if (comms[i] == nullptr) { + comms[i] = new RcclComm(); + } + SetRcclComm(comms[i], nccl_comms[static_cast(i)]); + } +} + +void RcclImpl::CommInitRank(CclComm **comm, int nranks, const CclUniqueId &unique_id, int rank) const { + CHECK_NOTNULL(comm); + CHECK_GT(nranks, 0); + + if (*comm == nullptr) { + *comm = new RcclComm(); + } + + ncclComm_t nccl_comm = nullptr; + NCCL_CHECK(ncclCommInitRank(&nccl_comm, nranks, GetRcclUniqueId(unique_id), rank)); + SetRcclComm(*comm, nccl_comm); +} + +void RcclImpl::CommDestroy(CclComm *comm) const { + if (comm == nullptr) { + return; + } + NCCL_CHECK(ncclCommDestroy(GetRcclComm(comm))); + SetRcclComm(comm, nullptr); +} + +void RcclImpl::AllReduce(const void *sendbuff, void *recvbuff, size_t count, DataType dtype, + nn::parallel::function::ReduceOpType reduce_op, const CclComm *comm, Stream *stream) const { + NCCL_CHECK(ncclAllReduce(sendbuff, recvbuff, count, kRcclDtypeMap.at(dtype), kRcclReduceOpMap.at(reduce_op), + GetRcclComm(comm), GetDcuStream(stream))); +} + +void RcclImpl::Broadcast(const void *sendbuff, void *recvbuff, size_t count, DataType dtype, int root, + const CclComm *comm, Stream *stream) const { + NCCL_CHECK(ncclBroadcast(sendbuff, recvbuff, count, kRcclDtypeMap.at(dtype), root, GetRcclComm(comm), + GetDcuStream(stream))); +} + +void RcclImpl::Reduce(const void *sendbuff, void *recvbuff, size_t count, DataType dtype, + nn::parallel::function::ReduceOpType reduce_op, int root, const CclComm *comm, + Stream *stream) const { + NCCL_CHECK(ncclReduce(sendbuff, recvbuff, count, kRcclDtypeMap.at(dtype), kRcclReduceOpMap.at(reduce_op), root, + GetRcclComm(comm), GetDcuStream(stream))); +} + +void RcclImpl::AllGather(const void *sendbuff, void *recvbuff, size_t count, DataType dtype, const CclComm *comm, + Stream *stream) const { + NCCL_CHECK( + ncclAllGather(sendbuff, recvbuff, count, kRcclDtypeMap.at(dtype), GetRcclComm(comm), GetDcuStream(stream))); +} + +void RcclImpl::ReduceScatter(const void *sendbuff, void *recvbuff, size_t recv_count, DataType dtype, + nn::parallel::function::ReduceOpType reduce_op, const CclComm *comm, + Stream *stream) const { + NCCL_CHECK(ncclReduceScatter(sendbuff, recvbuff, recv_count, kRcclDtypeMap.at(dtype), + kRcclReduceOpMap.at(reduce_op), GetRcclComm(comm), GetDcuStream(stream))); +} + +void RcclImpl::Send(const void *buff, size_t count, DataType dtype, int peer, const CclComm *comm, + Stream *stream) const { + NCCL_CHECK(ncclSend(buff, count, kRcclDtypeMap.at(dtype), peer, GetRcclComm(comm), GetDcuStream(stream))); +} + +void RcclImpl::Recv(void *buff, size_t count, DataType dtype, int peer, const CclComm *comm, Stream *stream) const { + NCCL_CHECK(ncclRecv(buff, count, kRcclDtypeMap.at(dtype), peer, GetRcclComm(comm), GetDcuStream(stream))); +} + +INFINI_TRAIN_REGISTER_CCL_IMPL(Device::DeviceType::kDCU, RcclImpl) + +} // namespace infini_train::core::dcu diff --git a/infini_train/src/core/ccl/dcu/rccl_impl.h b/infini_train/src/core/ccl/dcu/rccl_impl.h new file mode 100644 index 00000000..532f46e8 --- /dev/null +++ b/infini_train/src/core/ccl/dcu/rccl_impl.h @@ -0,0 +1,51 @@ +#pragma once + +#include +#include + +#include "infini_train/include/core/ccl/ccl.h" + +namespace infini_train::core::dcu { + +class RcclImpl final : public CclImpl { +public: + Device::DeviceType Type() const override; + + void GroupStart() const override; + + void GroupEnd() const override; + + void GetAsyncError(const CclComm *comm, CclStatus *async_error) const override; + + void GetUniqueId(CclUniqueId **unique_id) const override; + + void CommInitAll(CclComm **comms, int ndev, const int *devlist) const override; + + void CommInitRank(CclComm **comm, int nranks, const CclUniqueId &unique_id, int rank) const override; + + void CommDestroy(CclComm *comm) const override; + + void AllReduce(const void *sendbuff, void *recvbuff, size_t count, DataType dtype, + nn::parallel::function::ReduceOpType reduce_op, const CclComm *comm, Stream *stream) const override; + + void Broadcast(const void *sendbuff, void *recvbuff, size_t count, DataType dtype, int root, const CclComm *comm, + Stream *stream) const override; + + void Reduce(const void *sendbuff, void *recvbuff, size_t count, DataType dtype, + nn::parallel::function::ReduceOpType reduce_op, int root, const CclComm *comm, + Stream *stream) const override; + + void AllGather(const void *sendbuff, void *recvbuff, size_t count, DataType dtype, const CclComm *comm, + Stream *stream) const override; + + void ReduceScatter(const void *sendbuff, void *recvbuff, size_t recv_count, DataType dtype, + nn::parallel::function::ReduceOpType reduce_op, const CclComm *comm, + Stream *stream) const override; + + void Send(const void *buff, size_t count, DataType dtype, int peer, const CclComm *comm, + Stream *stream) const override; + + void Recv(void *buff, size_t count, DataType dtype, int peer, const CclComm *comm, Stream *stream) const override; +}; + +} // namespace infini_train::core::dcu diff --git a/infini_train/src/core/runtime/cpu/cpu_guard_impl.cc b/infini_train/src/core/runtime/cpu/cpu_guard_impl.cc index c298a413..16440b7e 100644 --- a/infini_train/src/core/runtime/cpu/cpu_guard_impl.cc +++ b/infini_train/src/core/runtime/cpu/cpu_guard_impl.cc @@ -2,7 +2,6 @@ #include #include -#include #include #include "glog/logging.h" @@ -133,9 +132,8 @@ void CpuGuardImpl::FreeAsync(void *dev_ptr, Stream *stream) { } void CpuGuardImpl::Memcpy(void *dst, const void *src, size_t count, MemcpyKind kind) { - CHECK(kind == MemcpyKind::kD2D) << std::format("CpuGuardImpl::Memcpy only supports kD2D (host-to-host) memcpy, " - "but got MemcpyKind={}", - MemcpyKindToString(kind)); + CHECK(kind == MemcpyKind::kD2D) << "CpuGuardImpl::Memcpy only supports kD2D (host-to-host) memcpy, but got " + << MemcpyKindToString(kind); std::memcpy(dst, src, count); } diff --git a/infini_train/src/core/runtime/dcu/dcu_dispatch.h b/infini_train/src/core/runtime/dcu/dcu_dispatch.h new file mode 100644 index 00000000..8bc2767f --- /dev/null +++ b/infini_train/src/core/runtime/dcu/dcu_dispatch.h @@ -0,0 +1,51 @@ +#pragma once + +#include +#include + +#include +#include + +#include "infini_train/include/core/backend_type_map.h" +#include "infini_train/include/dtype_dispatch.h" + +// ----------------------------------------------------------------------------- +// HIP low-precision BackendTypeMap specializations: +// FP16 -> __half, BF16 -> hip_bfloat16 +// ----------------------------------------------------------------------------- +namespace infini_train::core { +template <> struct BackendTypeMap { + using type = __half; +}; + +template <> struct BackendTypeMap { + using type = hip_bfloat16; +}; +} // namespace infini_train::core + +// Register all standard (non-low-precision) dtypes for the HIP backend. +// FP16/BF16 are registered explicitly above with their HIP-native scalar types. +INFINI_REGISTER_STANDARD_BACKEND_TYPES(infini_train::Device::DeviceType::kDCU) + +namespace infini_train::core::dcu { + +template struct DcuTypeMap : BackendTypeMap {}; + +// ----------------------------------------------------------------------------- +// HIP dispatch helpers +// ----------------------------------------------------------------------------- + +template +auto DispatchDcuFunc(DataType dtype, Functor &&func, std::string_view context_identifier = "", Args &&...args) { + return infini_train::DispatchByTypeMap( + dtype, std::forward(func), context_identifier, std::forward(args)...); +} + +template +auto DispatchDcuFunc(const std::vector &dtypes, Functor &&func, std::string_view context_identifier = "", + Args &&...args) { + return infini_train::DispatchByTypeMap( + dtypes, std::forward(func), context_identifier, std::forward(args)...); +} + +} // namespace infini_train::core::dcu diff --git a/infini_train/src/core/runtime/dcu/dcu_guard_impl.cc b/infini_train/src/core/runtime/dcu/dcu_guard_impl.cc new file mode 100644 index 00000000..534dccfc --- /dev/null +++ b/infini_train/src/core/runtime/dcu/dcu_guard_impl.cc @@ -0,0 +1,303 @@ +#include "infini_train/src/core/runtime/dcu/dcu_guard_impl.h" + +#include +#include +#include + +#include "infini_train/include/common/dcu/common_dcu.h" +#include "infini_train/include/core/runtime/runtime_common.h" +#include "infini_train/include/device.h" + +#include "infini_train/src/core/runtime/dcu/dcu_runtime_common.h" + +namespace infini_train::core::dcu { +namespace { +constexpr int kMaxGpus = 8; +constexpr size_t kBytesPerMB = 1024ULL * 1024ULL; + +static std::array, kMaxGpus> hip_streams; +static std::array, kMaxGpus> hip_blas_handles; + +static std::array device_stream_flags; +static std::array device_handle_flags; + +inline void CheckDcuDevice(Device device) { + CHECK(device.type() == Device::DeviceType::kDCU) + << "DcuGuardImpl expects HIP device, but got type=" << static_cast(device.type()) + << " index=" << static_cast(device.index()); + const int idx = device.index(); + CHECK(idx >= 0 && idx < kMaxGpus) << "HIP device index " << idx << " out of cache range [0, " << kMaxGpus << ")."; +} + +inline hipEvent_t GetDcuEvent(Event *event) { + auto *hip_event = dynamic_cast(event); + CHECK_NOTNULL(hip_event); + return hip_event->hip_event(); +} + +inline hipStream_t GetDcuStream(Stream *stream) { + auto *hip_stream = dynamic_cast(stream); + CHECK_NOTNULL(hip_stream); + return hip_stream->hip_stream(); +} +} // namespace + +void DcuGuardImpl::InitSingleStream(Device device) { + CheckDcuDevice(device); + + int current_device = -1; + HIP_CHECK(hipGetDevice(¤t_device)); + HIP_CHECK(hipSetDevice(device.index())); + + hip_streams[device.index()] = std::make_unique(); + + HIP_CHECK(hipSetDevice(current_device)); +} + +void DcuGuardImpl::InitSingleHandle(Device device) { + CheckDcuDevice(device); + + int current_device = -1; + HIP_CHECK(hipGetDevice(¤t_device)); + HIP_CHECK(hipSetDevice(device.index())); + + std::call_once(device_stream_flags.at(device.index()), InitSingleStream, device); + + hip_blas_handles[device.index()] = std::make_unique(hip_streams[device.index()].get()); + + HIP_CHECK(hipSetDevice(current_device)); +} + +DcuGuardImpl::DcuGuardImpl() {} + +// device +Device DcuGuardImpl::GetDevice() const { + int current_device = -1; + HIP_CHECK(hipGetDevice(¤t_device)); + return Device(Device::DeviceType::kDCU, current_device); +} + +void DcuGuardImpl::SetDevice(Device device) const { + CheckDcuDevice(device); + HIP_CHECK(hipSetDevice(device.index())); +} + +int DcuGuardImpl::DeviceCount() const { + int device_count = 0; + HIP_CHECK(hipGetDeviceCount(&device_count)); + return device_count; +} + +Device::DeviceType DcuGuardImpl::Type() const { return Device::DeviceType::kDCU; } + +// stream +Stream *DcuGuardImpl::GetStream(Device device) const { + CheckDcuDevice(device); + // FIXME(dcj): call_once is process-scoped and assumes single initialization. + // This can be problematic if the HIP backend is initialized multiple + // times within the same process (e.g. in unit tests). + std::call_once(device_stream_flags.at(device.index()), InitSingleStream, device); + return hip_streams.at(device.index()).get(); +} + +Stream *DcuGuardImpl::CreateStream(Device device) const { + CheckDcuDevice(device); + int current_device = -1; + HIP_CHECK(hipGetDevice(¤t_device)); + HIP_CHECK(hipSetDevice(device.index())); + + Stream *stream = new DcuStream(); + + HIP_CHECK(hipSetDevice(current_device)); + return stream; +} + +Stream *DcuGuardImpl::CreateStreamWithPriority(Device device, int priority) const { + CheckDcuDevice(device); + int current_device = -1; + HIP_CHECK(hipGetDevice(¤t_device)); + HIP_CHECK(hipSetDevice(device.index())); + + Stream *stream = new DcuStream(priority); + + HIP_CHECK(hipSetDevice(current_device)); + return stream; +} + +void DcuGuardImpl::DestroyStream(Stream *stream) const { + if (stream == nullptr) { + return; + } + auto *hip_stream = dynamic_cast(stream); + CHECK_NOTNULL(hip_stream); + delete hip_stream; +} + +void DcuGuardImpl::GetStreamPriorityRange(int *low, int *high) const { + HIP_CHECK(hipDeviceGetStreamPriorityRange(low, high)); +} + +// event +void DcuGuardImpl::EventCreate(Event **event) const { *event = new DcuEvent(); } + +void DcuGuardImpl::EventCreateWithFlags(Event **event, EventFlag flags) const { *event = new DcuEvent(flags); } + +void DcuGuardImpl::EventDestroy(Event *event) const { + if (event == nullptr) { + return; + } + delete event; +} + +void DcuGuardImpl::EventRecord(Event *event, Stream *stream) const { + auto hip_event = GetDcuEvent(event); + auto hip_stream = GetDcuStream(stream); + HIP_CHECK(hipEventRecord(hip_event, hip_stream)); +} + +void DcuGuardImpl::StreamWaitEvent(Stream *stream, Event *event, uint32_t flags) const { + auto hip_event = GetDcuEvent(event); + auto hip_stream = GetDcuStream(stream); + HIP_CHECK(hipStreamWaitEvent(hip_stream, hip_event, flags)); +} + +RuntimeStatus DcuGuardImpl::EventSynchronize(Event *event) const { + auto hip_event = GetDcuEvent(event); + hipError_t status = hipEventSynchronize(hip_event); + if (status == hipSuccess) { + return RuntimeStatus::kSuccess; + } + if (status == hipErrorNotReady) { + return RuntimeStatus::kNotReady; + } + LOG(ERROR) << "DcuGuardImpl::EventSynchronize failed: " << hipGetErrorString(status); + return RuntimeStatus::kError; +} + +RuntimeStatus DcuGuardImpl::EventQuery(Event *event) const { + auto hip_event = GetDcuEvent(event); + hipError_t status = hipEventQuery(hip_event); + if (status == hipSuccess) { + return RuntimeStatus::kSuccess; + } + if (status == hipErrorNotReady) { + return RuntimeStatus::kNotReady; + } + LOG(ERROR) << "DcuGuardImpl::EventQuery failed: " << hipGetErrorString(status); + return RuntimeStatus::kError; +} + +float DcuGuardImpl::EventElapsedTime(Event *start_event, Event *stop_event) const { + auto start_hip_event = GetDcuEvent(start_event); + auto stop_hip_event = GetDcuEvent(stop_event); + float elapsed_ms = 0.0f; + HIP_CHECK(hipEventElapsedTime(&elapsed_ms, start_hip_event, stop_hip_event)); + return elapsed_ms; +} + +// sync +void DcuGuardImpl::SynchronizeDevice(Device device) const { + auto original_device = GetDevice(); + SetDevice(device); + + HIP_CHECK(hipDeviceSynchronize()); + + SetDevice(original_device); +} + +void DcuGuardImpl::SynchronizeStream(Stream *stream) const { + auto hip_stream = GetDcuStream(stream); + HIP_CHECK(hipStreamSynchronize(hip_stream)); +} + +// blas +BlasHandle *DcuGuardImpl::GetBlasHandle(Device device) const { + CheckDcuDevice(device); + std::call_once(device_handle_flags.at(device.index()), InitSingleHandle, device); + return hip_blas_handles.at(device.index()).get(); +} + +// memory +void DcuGuardImpl::Malloc(void **dev_ptr, size_t size) { HIP_CHECK(hipMalloc(dev_ptr, size)); } + +void DcuGuardImpl::MallocAsync(void **dev_ptr, size_t size, Stream *stream) { + auto hip_stream = GetDcuStream(stream); + HIP_CHECK(hipMallocAsync(dev_ptr, size, hip_stream)); +} + +void DcuGuardImpl::Free(void *dev_ptr) { HIP_CHECK(hipFree(dev_ptr)); } + +void DcuGuardImpl::FreeAsync(void *dev_ptr, Stream *stream) { + auto hip_stream = GetDcuStream(stream); + HIP_CHECK(hipFreeAsync(dev_ptr, hip_stream)); +} + +void DcuGuardImpl::Memcpy(void *dst, const void *src, size_t count, MemcpyKind kind) { + if (kind == MemcpyKind::kH2D) { + HIP_CHECK(hipMemcpy(dst, src, count, hipMemcpyHostToDevice)); + } else if (kind == MemcpyKind::kD2H) { + HIP_CHECK(hipMemcpy(dst, src, count, hipMemcpyDeviceToHost)); + } else if (kind == MemcpyKind::kD2D) { + HIP_CHECK(hipMemcpy(dst, src, count, hipMemcpyDeviceToDevice)); + } else { + LOG(FATAL) << "DcuGuardImpl::Memcpy got invalid MemcpyKind=" << MemcpyKindToString(kind); + } +} + +void DcuGuardImpl::MemcpyAsync(void *dst, const void *src, size_t count, MemcpyKind kind, Stream *stream) { + auto hip_stream = GetDcuStream(stream); + + switch (kind) { + case MemcpyKind::kH2D: + HIP_CHECK(hipMemcpyAsync(dst, src, count, hipMemcpyHostToDevice, hip_stream)); + break; + case MemcpyKind::kD2H: + HIP_CHECK(hipMemcpyAsync(dst, src, count, hipMemcpyDeviceToHost, hip_stream)); + break; + case MemcpyKind::kD2D: + HIP_CHECK(hipMemcpyAsync(dst, src, count, hipMemcpyDeviceToDevice, hip_stream)); + break; + default: + LOG(FATAL) << "DcuGuardImpl::MemcpyAsync got invalid MemcpyKind=" << MemcpyKindToString(kind); + } +} + +void DcuGuardImpl::ResetMemPoolHighWatermarks(Device device) const { + int current_device = -1; + HIP_CHECK(hipGetDevice(¤t_device)); + + SetDevice(device); + hipMemPool_t pool; + HIP_CHECK(hipDeviceGetDefaultMemPool(&pool, device.index())); + + uint64_t zero = 0; + // High watermark can only be reset to zero; non-zero is illegal. + HIP_CHECK(hipMemPoolSetAttribute(pool, hipMemPoolAttrUsedMemHigh, &zero)); + HIP_CHECK(hipMemPoolSetAttribute(pool, hipMemPoolAttrReservedMemHigh, &zero)); + + HIP_CHECK(hipSetDevice(current_device)); +} + +std::pair DcuGuardImpl::GetMemPoolPeakMB(Device device) const { + int current_device = -1; + HIP_CHECK(hipGetDevice(¤t_device)); + + SetDevice(device); + hipMemPool_t pool; + HIP_CHECK(hipDeviceGetDefaultMemPool(&pool, device.index())); + + uint64_t used = 0; + HIP_CHECK(hipMemPoolGetAttribute(pool, hipMemPoolAttrUsedMemHigh, &used)); + + uint64_t reserved = 0; + HIP_CHECK(hipMemPoolGetAttribute(pool, hipMemPoolAttrReservedMemHigh, &reserved)); + + HIP_CHECK(hipSetDevice(current_device)); + + return std::make_pair(static_cast(used / kBytesPerMB), + static_cast(reserved / kBytesPerMB)); +} + +INFINI_TRAIN_REGISTER_DEVICE_GUARD_IMPL(Device::DeviceType::kDCU, DcuGuardImpl) + +} // namespace infini_train::core::dcu diff --git a/infini_train/src/core/runtime/dcu/dcu_guard_impl.h b/infini_train/src/core/runtime/dcu/dcu_guard_impl.h new file mode 100644 index 00000000..e064194e --- /dev/null +++ b/infini_train/src/core/runtime/dcu/dcu_guard_impl.h @@ -0,0 +1,85 @@ +#pragma once + +#include "infini_train/include/core/runtime/device_guard.h" +#include "infini_train/include/device.h" + +namespace infini_train::core { +class Stream; +class BlasHandle; +} // namespace infini_train::core + +namespace infini_train::core::dcu { + +class DcuGuardImpl final : public DeviceGuardImpl { +public: + static void InitSingleStream(Device device); + + static void InitSingleHandle(Device device); + + DcuGuardImpl(); + + // device + Device GetDevice() const override; + + void SetDevice(Device device) const override; + + int DeviceCount() const override; + + Device::DeviceType Type() const override; + + // stream + // TODO(zbl): Better wrap the create/destroy API call inside the constructor/destructor of class 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 + // TODO(zbl): Better wrap the create/destroy API call inside the constructor/destructor of class 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; +}; + +} // namespace infini_train::core::dcu diff --git a/infini_train/src/core/runtime/dcu/dcu_runtime_common.cc b/infini_train/src/core/runtime/dcu/dcu_runtime_common.cc new file mode 100644 index 00000000..273bf9fd --- /dev/null +++ b/infini_train/src/core/runtime/dcu/dcu_runtime_common.cc @@ -0,0 +1,58 @@ +#include "infini_train/src/core/runtime/dcu/dcu_runtime_common.h" + +#include "infini_train/include/common/dcu/common_dcu.h" + +namespace infini_train::core::dcu { +namespace { +uint32_t ToDcuEventFlags(EventFlag flags) { + switch (flags) { + case EventFlag::kDefault: + return hipEventDefault; + case EventFlag::kBlockingSync: + return hipEventBlockingSync; + case EventFlag::kDisableTiming: + return hipEventDisableTiming; + case EventFlag::kInterprocess: + // HIP requires hipEventDisableTiming with interprocess events. + return hipEventInterprocess | hipEventDisableTiming; + default: + LOG(FATAL) << "Unsupported EventFlag value: " << static_cast(flags); + } + return hipEventDefault; +} +} // namespace + +DcuEvent::DcuEvent(EventFlag flags) { HIP_CHECK(hipEventCreateWithFlags(&event_, ToDcuEventFlags(flags))); } + +DcuEvent::~DcuEvent() { + if (event_ != nullptr) { + HIP_CHECK(hipEventDestroy(event_)); + } +} + +hipEvent_t DcuEvent::hip_event() const { return event_; } + +DcuStream::DcuStream() { HIP_CHECK(hipStreamCreate(&stream_)); } + +DcuStream::DcuStream(int priority) { + HIP_CHECK(hipStreamCreateWithPriority(&stream_, hipStreamNonBlocking, priority)); +} + +DcuStream::~DcuStream() { + // Do nothing. +} + +hipStream_t DcuStream::hip_stream() const { return stream_; } + +DcuBlasHandle::DcuBlasHandle(Stream *stream) { + HIPBLAS_CHECK(hipblasCreate(&hipblas_handle_)); + HIPBLAS_CHECK(hipblasSetStream(hipblas_handle_, dynamic_cast(stream)->hip_stream())); +} + +DcuBlasHandle::~DcuBlasHandle() { + // Do nothing. +} + +hipblasHandle_t DcuBlasHandle::hipblas_handle() const { return hipblas_handle_; } + +} // namespace infini_train::core::dcu diff --git a/infini_train/src/core/runtime/dcu/dcu_runtime_common.h b/infini_train/src/core/runtime/dcu/dcu_runtime_common.h new file mode 100644 index 00000000..4d4ad093 --- /dev/null +++ b/infini_train/src/core/runtime/dcu/dcu_runtime_common.h @@ -0,0 +1,64 @@ +#pragma once + +#include + +#include +#include + +#include "infini_train/include/core/runtime/runtime_common.h" + +namespace infini_train::core { +class Stream; +} + +namespace infini_train::core::dcu { + +class DcuEvent final : public Event { +public: + explicit DcuEvent(EventFlag flags = EventFlag::kDefault); + ~DcuEvent() override; + + hipEvent_t hip_event() const; + +private: + hipEvent_t event_ = nullptr; +}; + +class DcuStream : public Stream { +public: + DcuStream(); + explicit DcuStream(int priority); + + // NOTE(dcj): + // The DcuStream are "leaked": they are created but never destroyed because the + // destruction of global variables could happen after the HIP runtime has + // already been destroyed and thus invoking hipStreamDestroy could lead to a + // crash. It's likely an issue in HIP, but to be safe - let's just "forget" + // the destruction. + ~DcuStream() override; + + hipStream_t hip_stream() const; + +private: + hipStream_t stream_ = nullptr; +}; + +class DcuBlasHandle : public BlasHandle { +public: + explicit DcuBlasHandle(Stream *stream); + + // NOTE(dcj): + // The DcuBlasHandle are "leaked": they are created but never destroyed because the + // destruction of global variables could happen after the HIP runtime has + // already been destroyed and thus invoking chipblasDestroy could lead to a + // crash. It's likely an issue in HIP, but to be safe - let's just "forget" + // the destruction. + ~DcuBlasHandle() override; + + hipblasHandle_t hipblas_handle() const; + +private: + hipblasHandle_t hipblas_handle_; +}; + +} // namespace infini_train::core::dcu diff --git a/infini_train/src/core/runtime/device_guard.cc b/infini_train/src/core/runtime/device_guard.cc index fbcb316f..24d560e3 100644 --- a/infini_train/src/core/runtime/device_guard.cc +++ b/infini_train/src/core/runtime/device_guard.cc @@ -1,6 +1,5 @@ #include "infini_train/include/core/runtime/device_guard.h" -#include #include #include @@ -135,19 +134,19 @@ DeviceGuardImplRegistry &DeviceGuardImplRegistry::Instance() { void DeviceGuardImplRegistry::Register(Device::DeviceType type, std::unique_ptr impl) { if (type != impl->Type()) { - LOG(FATAL) << std::format("Register device guard impl with type {}, but as type {}", - static_cast(impl->Type()), static_cast(type)); + LOG(FATAL) << "Register device guard impl with type " << static_cast(impl->Type()) << ", but as type " + << static_cast(type); } if (impls_.contains(type)) { - LOG(FATAL) << std::format("DeviceGuardImpl for type {} already registrered", static_cast(type)); + LOG(FATAL) << "DeviceGuardImpl for type " << static_cast(type) << " already registrered"; } if (!impls_.empty()) { for (auto &kv : impls_) { if (kv.first != Device::DeviceType::kCPU) { - LOG(FATAL) << std::format("Only CPU and one GPU backend allowed. Already have GPU={}, new={} rejected.", - static_cast(kv.first), static_cast(type)); + LOG(FATAL) << "Only CPU and one GPU backend allowed. Already have GPU=" << static_cast(kv.first) + << ", new=" << static_cast(type) << " rejected."; } } } diff --git a/infini_train/src/device.cc b/infini_train/src/device.cc index 1bb3aaad..296969f3 100644 --- a/infini_train/src/device.cc +++ b/infini_train/src/device.cc @@ -1,7 +1,6 @@ #include "infini_train/include/device.h" #include -#include #include #include @@ -26,9 +25,36 @@ bool Device::IsCPU() const { return type_ == DeviceType::kCPU; } bool Device::IsCUDA() const { return type_ == DeviceType::kCUDA; } +<<<<<<< ours std::string Device::ToString() const { std::ostringstream oss; oss << std::format("Device({}, {})", type_ == DeviceType::kCPU ? "CPU" : "CUDA", index_); +======= +bool Device::IsMACA() const { return type_ == DeviceType::kMACA; } + +bool Device::IsDCU() const { return type_ == DeviceType::kDCU; } + +std::string Device::ToString() const { + const char *type_str = "Unknown"; + switch (type_) { + case DeviceType::kCPU: + type_str = "CPU"; + break; + case DeviceType::kCUDA: + type_str = "CUDA"; + break; + case DeviceType::kMACA: + type_str = "MACA"; + break; + case DeviceType::kDCU: + type_str = "DCU"; + break; + default: + break; + } + std::ostringstream oss; + oss << "Device(" << type_str << ", " << static_cast(index_) << ")"; +>>>>>>> theirs return oss.str(); } diff --git a/infini_train/src/kernels/dcu/accumulate_grad.hip b/infini_train/src/kernels/dcu/accumulate_grad.hip new file mode 100644 index 00000000..a19cea7d --- /dev/null +++ b/infini_train/src/kernels/dcu/accumulate_grad.hip @@ -0,0 +1,95 @@ +#include +#include + +#include "infini_train/include/common/dcu/kernel_helper.cuh" +#include "infini_train/include/core/runtime/device_guard.h" +#include "infini_train/include/dispatcher.h" +#include "infini_train/include/tensor.h" + +#include "infini_train/src/core/runtime/dcu/dcu_dispatch.h" +#include "infini_train/src/core/runtime/dcu/dcu_runtime_common.h" + +namespace infini_train::kernels::dcu { + +template +__global__ void AccumulateGradKernel(const T *grad_ptr, float rate, T *tensor_ptr, size_t num_elements) { + int idx = blockIdx.x * blockDim.x + threadIdx.x; + if (idx < num_elements) { + tensor_ptr[idx] += common::dcu::Mul(grad_ptr[idx], common::dcu::Cast(rate)); + } +} + +void AccumulateGrad(const std::shared_ptr &gradient, float rate, const std::shared_ptr &tensor) { + size_t num_elements = gradient->NumElements(); + + int threads_per_block = 256; + int num_blocks = (num_elements + threads_per_block - 1) / threads_per_block; + + auto device = tensor->GetDevice(); + const auto &hip_stream = dynamic_cast( + infini_train::core::GetDeviceGuardImpl(device.type())->GetStream(device)) + ->hip_stream(); + + core::dcu::DispatchDcuFunc( + gradient->Dtype(), + [=]() { + AccumulateGradKernel<<>>( + static_cast(gradient->DataPtr()), rate, static_cast(tensor->DataPtr()), num_elements); + }, + "HIP AccumulateGrad"); +} + +template +__global__ void AdamAccumulateGradKernel(const T *grad_data, T *param_data, size_t num_elements, T *m_data, T *v_data, + float learning_rate, float beta1, float beta2, float eps, + const float bias_correction_m, const float bias_correction_v) { + size_t idx = blockIdx.x * blockDim.x + threadIdx.x; + if (idx < num_elements) { + m_data[idx] = common::dcu::Fma(common::dcu::Cast(beta1), m_data[idx], + common::dcu::Cast(1 - beta1) * grad_data[idx]); + v_data[idx] = common::dcu::Fma(common::dcu::Cast(beta2), v_data[idx], + common::dcu::Cast(1 - beta2) * grad_data[idx] * grad_data[idx]); + + const float m_hat = common::dcu::Cast(m_data[idx]) / bias_correction_m; + const float v_hat = common::dcu::Cast(v_data[idx]) / bias_correction_v; + + param_data[idx] = common::dcu::Sub( + param_data[idx], common::dcu::Cast(learning_rate * m_hat * __frcp_rn(__fsqrt_rn(v_hat) + eps))); + } +} + +void AdamAccumulateGrad(const std::shared_ptr &grad, const std::shared_ptr ¶m, + const std::shared_ptr &m, const std::shared_ptr &v, float learning_rate, + float beta1, float beta2, float eps, int64_t t) { + size_t num_elements = grad->NumElements(); + + const float bias_correction_m = 1.0f - std::pow(beta1, t); + const float bias_correction_v = 1.0f - std::pow(beta2, t); + + int threads_per_block = 256; + int num_blocks = (num_elements + threads_per_block - 1) / threads_per_block; + + auto device = grad->GetDevice(); + const auto &hip_stream = dynamic_cast( + infini_train::core::GetDeviceGuardImpl(device.type())->GetStream(device)) + ->hip_stream(); + + core::dcu::DispatchDcuFunc( + grad->Dtype(), + [=]() { + AdamAccumulateGradKernel<<>>( + static_cast(grad->DataPtr()), static_cast(param->DataPtr()), num_elements, + static_cast(m->DataPtr()), static_cast(v->DataPtr()), learning_rate, beta1, beta2, eps, + bias_correction_m, bias_correction_v); + }, + "HIP AdamAccumulateGrad"); +} +} // namespace infini_train::kernels::dcu + +#define REGISTER_HIP_ACCUMULATE_GRAD_KERNEL(kernel_name) \ + REGISTER_KERNEL(infini_train::Device::DeviceType::kDCU, kernel_name, infini_train::kernels::dcu::kernel_name) + +REGISTER_HIP_ACCUMULATE_GRAD_KERNEL(AccumulateGrad) +REGISTER_HIP_ACCUMULATE_GRAD_KERNEL(AdamAccumulateGrad) + +#undef REGISTER_HIP_ACCUMULATE_GRAD_KERNEL diff --git a/infini_train/src/kernels/dcu/cast.hip b/infini_train/src/kernels/dcu/cast.hip new file mode 100644 index 00000000..978f2a38 --- /dev/null +++ b/infini_train/src/kernels/dcu/cast.hip @@ -0,0 +1,58 @@ +#include + +#include "infini_train/include/common/common.h" +#include "infini_train/include/common/dcu/kernel_helper.cuh" +#include "infini_train/include/core/runtime/device_guard.h" +#include "infini_train/include/datatype.h" +#include "infini_train/include/device.h" +#include "infini_train/include/dispatcher.h" +#include "infini_train/include/tensor.h" + +#include "infini_train/src/core/runtime/dcu/dcu_dispatch.h" +#include "infini_train/src/core/runtime/dcu/dcu_runtime_common.h" + +namespace infini_train::kernels::dcu { + +template +__global__ void CastKernel(Tdst *dst, const Tsrc *src, size_t num_elements, size_t offset) { + size_t idx = blockIdx.x * blockDim.x + threadIdx.x + offset; + + if (idx < num_elements) { + dst[idx] = common::dcu::Cast(src[idx]); + } +} + +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(); + const auto &hip_stream = dynamic_cast( + infini_train::core::GetDeviceGuardImpl(device.type())->GetStream(device)) + ->hip_stream(); + + const size_t num_elements = input->NumElements(); + dim3 block_dims(256); + dim3 grid_dims(CEIL_DIV(num_elements, block_dims.x)); + const size_t step = grid_dims.x * block_dims.x; + + core::dcu::DispatchDcuFunc, DataTypeList>( + {dtype, input->Dtype()}, + [=]() { + auto dst = static_cast(dst_tensor->DataPtr()); + auto src = static_cast(input->DataPtr()); + + for (size_t offset = 0; offset < num_elements; offset += step) { + CastKernel<<>>(dst, src, num_elements, offset); + } + }, + "HIP Cast"); + + return {dst_tensor}; +} +} // namespace infini_train::kernels::dcu + +#define REGISTER_HIP_CAST_KERNEL(kernel_name) \ + REGISTER_KERNEL(infini_train::Device::DeviceType::kDCU, kernel_name, infini_train::kernels::dcu::kernel_name) + +REGISTER_HIP_CAST_KERNEL(Cast) + +#undef REGISTER_HIP_CAST_KERNEL diff --git a/infini_train/src/kernels/dcu/comm.hip b/infini_train/src/kernels/dcu/comm.hip new file mode 100644 index 00000000..79c47d99 --- /dev/null +++ b/infini_train/src/kernels/dcu/comm.hip @@ -0,0 +1,81 @@ +#include +#include +#include + +#include "glog/logging.h" + +#include "infini_train/include/device.h" +#include "infini_train/include/dispatcher.h" +#include "infini_train/include/nn/functional.h" +#include "infini_train/include/tensor.h" + +namespace infini_train::kernels::dcu { + +std::vector> Broadcast(const std::vector> &input_tensors, + const std::vector &devices) { + std::vector> outputs; + for (int i = 0; i < devices.size(); ++i) { + for (const auto &tensor : input_tensors) { + outputs.push_back(std::make_shared(tensor->To(devices[i]))); + } + } + return outputs; +} + +std::vector> ReduceAddCoalesced(const std::vector>> &grads, + Device destination) { + std::vector> outputs; + auto kernel = Dispatcher::Instance().GetKernel({destination.type(), "AccumulateGrad"}); + std::vector>> to_destination_grads; + for (int i = 0; i < grads[0].size(); ++i) { + outputs.emplace_back(std::make_shared(grads[0][i]->Dims(), grads[0][i]->Dtype(), destination)); + outputs[i]->Fill(0.0); + } + for (int i = 0; i < grads.size(); ++i) { + to_destination_grads.push_back(std::vector>()); + for (int j = 0; j < grads[i].size(); ++j) { + to_destination_grads[i].push_back(std::make_shared(grads[i][j]->To(destination))); + } + } + for (int i = 0; i < grads.size(); ++i) { + for (int j = 0; j < grads[i].size(); ++j) { + kernel.Call(to_destination_grads[i][j], static_cast(1.0), outputs[j]); + } + } + return outputs; +} + +std::vector> Scatter(const std::shared_ptr &tensor, std::vector devices, + int64_t dim) { + std::vector> outputs; + // FIXME(dcj): do split without autograd + std::vector> split_tensors = tensor->Split(tensor->Dims()[dim] / devices.size(), dim); + for (auto i = 0; i < devices.size(); ++i) { + outputs.push_back(std::make_shared(split_tensors[i]->To(devices[i]))); + } + return outputs; +} + +std::shared_ptr Gather(const std::vector> &tensors, Device destination, int64_t dim) { + std::vector> outputs; + for (const auto &tensor : tensors) { outputs.push_back(std::make_shared(tensor->To(destination))); } + auto kernel = Dispatcher::Instance().GetKernel({tensors[0]->GetDevice().type(), "StackForward"}); + auto gathered_tensor = kernel.Call>(outputs, dim); + auto old_dims = gathered_tensor->Dims(); + std::vector new_dims{old_dims[0] * old_dims[1]}; + for (int i = 2; i < old_dims.size(); ++i) { new_dims.push_back(old_dims[i]); } + auto view_kernel = Dispatcher::Instance().GetKernel({destination.type(), "NoOpForward"}); + return view_kernel.Call>(gathered_tensor, new_dims); +} +} // namespace infini_train::kernels::dcu + +#define REGISTER_HIP_COMM_KERNEL(kernel_name) \ + REGISTER_KERNEL(infini_train::Device::DeviceType::kDCU, Comm##kernel_name, \ + infini_train::kernels::dcu::kernel_name) + +REGISTER_HIP_COMM_KERNEL(Broadcast) +REGISTER_HIP_COMM_KERNEL(Scatter) +REGISTER_HIP_COMM_KERNEL(Gather) +REGISTER_HIP_COMM_KERNEL(ReduceAddCoalesced) + +#undef REGISTER_HIP_COMM_KERNEL diff --git a/infini_train/src/kernels/dcu/concat.hip b/infini_train/src/kernels/dcu/concat.hip new file mode 100644 index 00000000..d5880c84 --- /dev/null +++ b/infini_train/src/kernels/dcu/concat.hip @@ -0,0 +1,247 @@ +#include +#include +#include +#include +#include + +#include "glog/logging.h" + +#include "infini_train/include/common/dcu/common_dcu.h" +#include "infini_train/include/core/runtime/device_guard.h" +#include "infini_train/include/dispatcher.h" +#include "infini_train/include/tensor.h" + +#include "infini_train/src/core/runtime/dcu/dcu_dispatch.h" +#include "infini_train/src/core/runtime/dcu/dcu_runtime_common.h" + +namespace infini_train::kernels::dcu { +__device__ __forceinline__ int64_t UpperBoundI64(const int64_t *offsets, int64_t n_plus_1, int64_t x) { + // Return the largest s so that offsets[s] <= x + // offsets[0] = 0, offsets is monotonically increasing + // len(offsets) = num_inputs + 1 + int64_t l = 0, r = n_plus_1; // start search in [0, n+1) + while (l < r) { + int64_t m = l + ((r - l) >> 1); + if (offsets[m] <= x) { + l = m + 1; + } else { + r = m; + } + } + return l - 1; +} + +template +__global__ void ConcatForwardKernel(const T **inputs, T *output, const int64_t *offsets, int64_t N, int64_t D, + int64_t num_inputs, int64_t K_total) { + int64_t idx = blockIdx.x * blockDim.x + threadIdx.x; + int64_t total = N * K_total * D; + if (idx >= total) { + return; + } + + int64_t d = idx % D; + int64_t k = (idx / D) % K_total; + int64_t n = idx / (D * K_total); + + // find the largest s so that offsets[s] <= k < offsets[s+1] + int64_t s = UpperBoundI64(offsets, num_inputs + 1, k); + int64_t k_local = k - offsets[s]; + int64_t Ki = offsets[s + 1] - offsets[s]; + + const T *input = inputs[s]; + output[idx] = input[n * (Ki * D) + k_local * D + d]; +} + +std::shared_ptr ConcatForward(const std::vector> &inputs, int64_t dim) { + CHECK(!inputs.empty()); + + const auto &base_dims = inputs[0]->Dims(); + auto dtype = inputs[0]->Dtype(); + auto device = inputs[0]->GetDevice(); + + if (dim < 0) { + dim += static_cast(base_dims.size()); + } + CHECK_GE(dim, 0); + CHECK_LT(dim, static_cast(base_dims.size())); + + // Check shape requirements and save length along dim + std::vector Ks; + Ks.reserve(inputs.size()); + for (const auto &t : inputs) { + CHECK(t->Dtype() == dtype); + CHECK_EQ(t->Dims().size(), base_dims.size()); + for (size_t ax = 0; ax < base_dims.size(); ++ax) { + if (static_cast(ax) == dim) { + continue; + } + CHECK_EQ(t->Dims()[ax], base_dims[ax]) << "All non-concat dims must match"; + } + Ks.push_back(t->Dims()[dim]); + } + + std::vector out_dims = base_dims; + out_dims[dim] = std::accumulate(Ks.begin(), Ks.end(), int64_t{0}); + auto output = std::make_shared(out_dims, dtype, device); + + const int64_t N = std::accumulate(base_dims.begin(), base_dims.begin() + dim, 1LL, std::multiplies()); + const int64_t D = std::accumulate(base_dims.begin() + dim + 1, base_dims.end(), 1LL, std::multiplies()); + const int64_t num_inputs = static_cast(inputs.size()); + const int64_t K_total = out_dims[dim]; + + // offsets records the sum of Ks + // offsets[i] = sum_{j < i} K_j + std::vector host_offsets(num_inputs + 1, 0); + for (int64_t i = 0; i < num_inputs; ++i) { host_offsets[i + 1] = host_offsets[i] + Ks[i]; } + + const auto &stream = dynamic_cast( + infini_train::core::GetDeviceGuardImpl(device.type())->GetStream(device)) + ->hip_stream(); + + int64_t total = N * K_total * D; + int threads_per_block = 256; + int num_blocks = static_cast((total + threads_per_block - 1) / threads_per_block); + + core::dcu::DispatchDcuFunc( + dtype, + [=, &inputs, &host_offsets]() { + std::vector host_input_ptrs; + host_input_ptrs.reserve(inputs.size()); + for (const auto &t : inputs) { host_input_ptrs.push_back(static_cast(t->DataPtr())); } + + const T **device_input_ptrs = nullptr; + int64_t *device_offsets = nullptr; + + HIP_CHECK(hipMallocAsync(&device_input_ptrs, sizeof(T *) * num_inputs, stream)); + HIP_CHECK(hipMemcpyAsync(device_input_ptrs, host_input_ptrs.data(), sizeof(T *) * num_inputs, + hipMemcpyHostToDevice, stream)); + + HIP_CHECK(hipMallocAsync(&device_offsets, sizeof(int64_t) * (num_inputs + 1), stream)); + HIP_CHECK(hipMemcpyAsync(device_offsets, host_offsets.data(), sizeof(int64_t) * (num_inputs + 1), + hipMemcpyHostToDevice, stream)); + + ConcatForwardKernel<<>>( + device_input_ptrs, static_cast(output->DataPtr()), device_offsets, N, D, num_inputs, K_total); + + HIP_CHECK(hipFree(device_input_ptrs)); + HIP_CHECK(hipFree(device_offsets)); + }, + "HIP ConcatForward"); + + return output; +} + +template +__global__ void ConcatBackwardKernel(const T *grad_output, T **grad_inputs, const int64_t *offsets, int64_t N, + int64_t D, int64_t num_inputs, int64_t K_total) { + int64_t idx = blockIdx.x * blockDim.x + threadIdx.x; + int64_t total = N * K_total * D; + if (idx >= total) { + return; + } + + int64_t d = idx % D; + int64_t k = (idx / D) % K_total; + int64_t n = idx / (D * K_total); + + int64_t s = UpperBoundI64(offsets, num_inputs + 1, k); + int64_t k_local = k - offsets[s]; + int64_t Ki = offsets[s + 1] - offsets[s]; + + T *gi = grad_inputs[s]; + gi[n * (Ki * D) + k_local * D + d] = grad_output[idx]; +} + +std::vector> ConcatBackward(const std::shared_ptr &grad_output, + const std::vector> &input_dims_list, + int64_t dim) { + CHECK(!input_dims_list.empty()); + + auto dtype = grad_output->Dtype(); + const auto &output_dims = grad_output->Dims(); + if (dim < 0) { + dim += static_cast(output_dims.size()); + } + CHECK_GE(dim, 0); + CHECK_LT(dim, static_cast(output_dims.size())); + + const auto &base_rank = input_dims_list[0].size(); + std::vector Ks; + Ks.reserve(input_dims_list.size()); + for (const auto &dvec : input_dims_list) { + CHECK_EQ(dvec.size(), base_rank); + for (size_t ax = 0; ax < dvec.size(); ++ax) { + if (static_cast(ax) == dim) { + continue; + } + CHECK_EQ(dvec[ax], input_dims_list[0][ax]); + } + Ks.push_back(dvec[dim]); + } + + auto device = grad_output->GetDevice(); + + std::vector> grads; + grads.reserve(input_dims_list.size()); + for (const auto &dvec : input_dims_list) { + auto t = std::make_shared(dvec, dtype, device); + t->Fill(0.0); + grads.push_back(t); + } + + const int64_t N = std::accumulate(input_dims_list[0].begin(), input_dims_list[0].begin() + dim, 1LL, + std::multiplies()); + const int64_t D = std::accumulate(input_dims_list[0].begin() + dim + 1, input_dims_list[0].end(), 1LL, + std::multiplies()); + const int64_t num_inputs = static_cast(input_dims_list.size()); + const int64_t K_total = std::accumulate(Ks.begin(), Ks.end(), int64_t{0}); + + std::vector host_offsets(num_inputs + 1, 0); + for (int64_t i = 0; i < num_inputs; ++i) { host_offsets[i + 1] = host_offsets[i] + Ks[i]; } + + const auto &stream = dynamic_cast( + infini_train::core::GetDeviceGuardImpl(device.type())->GetStream(device)) + ->hip_stream(); + + int64_t total = N * K_total * D; + int threads_per_block = 256; + int num_blocks = static_cast((total + threads_per_block - 1) / threads_per_block); + + core::dcu::DispatchDcuFunc( + dtype, + [=, &grads, &host_offsets]() { + std::vector host_ptrs; + host_ptrs.reserve(grads.size()); + for (auto &t : grads) { host_ptrs.push_back(static_cast(t->DataPtr())); } + + T **device_ptrs = nullptr; + int64_t *device_offsets = nullptr; + + HIP_CHECK(hipMallocAsync(&device_ptrs, sizeof(T *) * num_inputs, stream)); + HIP_CHECK(hipMemcpyAsync(device_ptrs, host_ptrs.data(), sizeof(T *) * num_inputs, hipMemcpyHostToDevice, + stream)); + + HIP_CHECK(hipMallocAsync(&device_offsets, sizeof(int64_t) * (num_inputs + 1), stream)); + HIP_CHECK(hipMemcpyAsync(device_offsets, host_offsets.data(), sizeof(int64_t) * (num_inputs + 1), + hipMemcpyHostToDevice, stream)); + + ConcatBackwardKernel<<>>( + static_cast(grad_output->DataPtr()), device_ptrs, device_offsets, N, D, num_inputs, K_total); + + HIP_CHECK(hipFree(device_ptrs)); + HIP_CHECK(hipFree(device_offsets)); + }, + "HIP ConcatBackward"); + + return grads; +} +} // namespace infini_train::kernels::dcu + +#define REGISTER_HIP_CONCAT_KERNEL(kernel_name) \ + REGISTER_KERNEL(infini_train::Device::DeviceType::kDCU, kernel_name, infini_train::kernels::dcu::kernel_name) + +REGISTER_HIP_CONCAT_KERNEL(ConcatForward) +REGISTER_HIP_CONCAT_KERNEL(ConcatBackward) + +#undef REGISTER_HIP_CONCAT_KERNEL diff --git a/infini_train/src/kernels/dcu/cross_entropy.hip b/infini_train/src/kernels/dcu/cross_entropy.hip new file mode 100644 index 00000000..26b0dc43 --- /dev/null +++ b/infini_train/src/kernels/dcu/cross_entropy.hip @@ -0,0 +1,229 @@ +#include +#include +#include + +#include +#include + +#include "infini_train/include/common/dcu/common_dcu.h" +#include "infini_train/include/common/dcu/cub_compat.cuh" +#include "infini_train/include/common/dcu/kernel_helper.cuh" +#include "infini_train/include/core/runtime/device_guard.h" +#include "infini_train/include/dispatcher.h" +#include "infini_train/include/tensor.h" + +#include "infini_train/src/core/runtime/dcu/dcu_dispatch.h" +#include "infini_train/src/core/runtime/dcu/dcu_runtime_common.h" + +namespace infini_train::kernels::dcu { +namespace { +constexpr float kNegativeInfinity = -std::numeric_limits::infinity(); +} + +template +__global__ void CrossEntropyForwardKernel(const InputType *__restrict__ input_ptr, + const TargetType *__restrict__ target_ptr, InputType *__restrict__ loss_ptr, + int bs, int num_classes) { + __shared__ struct { + float max_logit; + float sum_exp; + TargetType target_class; + typename hipcub::BlockReduce::TempStorage reduce; + } shared; + + const int sample_idx = blockIdx.x; + if (sample_idx >= bs) { + return; + } + + const int tid = threadIdx.x; + const size_t base = sample_idx * num_classes; + + if (tid == 0) { + shared.target_class = target_ptr[sample_idx]; + } + __syncthreads(); + + // calculate the max + float thread_max = kNegativeInfinity; + for (int i = tid; i < num_classes; i += BLOCK_SIZE) { + thread_max = fmaxf(thread_max, common::dcu::Cast(input_ptr[base + i])); + } + const float block_max = hipcub::BlockReduce(shared.reduce).Reduce(thread_max, CubMaxOp()); + if (tid == 0) { + shared.max_logit = block_max; + } + __syncthreads(); + + // calculate the sum of exponents + float thread_sum = 0.0f; + for (int i = tid; i < num_classes; i += BLOCK_SIZE) { + thread_sum += expf(common::dcu::Cast(input_ptr[base + i]) - shared.max_logit); + } + const float block_sum = hipcub::BlockReduce(shared.reduce).Sum(thread_sum); + if (tid == 0) { + shared.sum_exp = block_sum; + } + __syncthreads(); + + // calculate the loss + if (tid == 0) { + const float target_val + = common::dcu::Cast(input_ptr[base + common::dcu::Cast(shared.target_class)]) + - shared.max_logit; + loss_ptr[sample_idx] = logf(shared.sum_exp) - target_val; + } +} + +std::shared_ptr CrossEntropyForward(const std::shared_ptr &input, + const std::shared_ptr &target) { + const auto &input_dims = input->Dims(); + CHECK_GE(input_dims.size(), 2); + const int bs = std::accumulate(input_dims.rbegin() + 1, input_dims.rend(), 1, std::multiplies{}); + const int num_classes = *input_dims.rbegin(); + + auto batched_output = std::make_shared(std::vector{bs}, input->Dtype(), input->GetDevice()); + + constexpr int threads_per_block = 256; + int num_blocks = bs; + + auto device = target->GetDevice(); + const auto &hip_stream = dynamic_cast( + infini_train::core::GetDeviceGuardImpl(device.type())->GetStream(device)) + ->hip_stream(); + + return core::dcu::DispatchDcuFunc, + DataTypeList>( + {target->Dtype(), input->Dtype()}, + [=]() { + const Ttarget *target_ptr = static_cast(target->DataPtr()); + const Tinput *input_ptr = static_cast(input->DataPtr()); + Tinput *batched_loss_ptr = static_cast(batched_output->DataPtr()); + // FIXME(dcj): do reduce on GPU + CrossEntropyForwardKernel + <<>>(input_ptr, target_ptr, batched_loss_ptr, bs, + num_classes); + + auto loss_cpu = batched_output->To(Device()); + auto loss = std::make_shared(std::vector{}, input->Dtype(), Device()); + auto loss_cpu_typed_ptr = static_cast(loss_cpu.DataPtr()); + static_cast(loss->DataPtr())[0] + = std::accumulate(loss_cpu_typed_ptr, loss_cpu_typed_ptr + bs, 0.0f, + [](float acc, const Tinput &val) { return acc + common::dcu::Cast(val); }) + / bs; + + return std::make_shared(loss->To(input->GetDevice())); + }, + "HIP CrossEntropyForward"); +} + +template +__global__ void CrossEntropyBackwardKernel(const InputType *__restrict__ input_ptr, + InputType *__restrict__ input_grad_ptr, + const TargetType *__restrict__ target_ptr, + const InputType *__restrict__ output_grad_ptr, int bs, int num_classes) { + __shared__ struct { + float max_logit; + float sum_exp; + int target_class; + typename hipcub::BlockReduce::TempStorage reduce; + } shared; + + const int tid = threadIdx.x; + const int idx = blockIdx.x; + + if (idx >= bs) { + return; + } + + const size_t idx_base = idx * num_classes; + + if (tid == 0) { + shared.target_class = static_cast(target_ptr[idx]); + } + __syncthreads(); + + // calculate the max + float thread_max = kNegativeInfinity; + for (int i = tid; i < num_classes; i += BLOCK_SIZE) { + thread_max = fmaxf(thread_max, common::dcu::Cast(input_ptr[idx_base + i])); + } + const float block_max = hipcub::BlockReduce(shared.reduce).Reduce(thread_max, CubMaxOp()); + if (tid == 0) { + shared.max_logit = block_max; + } + __syncthreads(); + + // calculate the sum + float thread_sum = 0.0f; + for (int i = tid; i < num_classes; i += BLOCK_SIZE) { + thread_sum += expf(common::dcu::Cast(input_ptr[idx_base + i]) - shared.max_logit); + } + + const float block_sum = hipcub::BlockReduce(shared.reduce).Sum(thread_sum); + if (tid == 0) { + shared.sum_exp = block_sum; + } + __syncthreads(); + + // calculate the gradient + const float inv_bs = 1.0f / bs; + const float scale = 1.0f / shared.sum_exp; + const int target = shared.target_class; + + for (int i = tid; i < num_classes; i += BLOCK_SIZE) { + const int global_idx = idx_base + i; + const float exp_val = expf(common::dcu::Cast(input_ptr[global_idx]) - shared.max_logit); + input_grad_ptr[global_idx] = common::dcu::Cast((exp_val * scale - (i == target)) * inv_bs + * common::dcu::Cast(output_grad_ptr[0])); + } +} + +std::shared_ptr CrossEntropyBackward(const std::shared_ptr &input, + const std::shared_ptr &target, + const std::shared_ptr &grad_output) { + const auto &input_dims = input->Dims(); + CHECK_GE(input_dims.size(), 2); + const int bs = std::accumulate(input_dims.rbegin() + 1, input_dims.rend(), 1, std::multiplies{}); + const int num_classes = *input_dims.rbegin(); + + auto input_casted = std::make_shared(input->To(grad_output->Dtype())); + + CHECK_EQ(grad_output->Dims().size(), 0); + auto grad_input = std::make_shared(input_casted->Dims(), input_casted->Dtype(), grad_output->GetDevice()); + + constexpr int threads_per_block = 256; + int num_blocks = bs; + + auto device = target->GetDevice(); + const auto &hip_stream = dynamic_cast( + infini_train::core::GetDeviceGuardImpl(device.type())->GetStream(device)) + ->hip_stream(); + + core::dcu::DispatchDcuFunc, + DataTypeList>( + {target->Dtype(), input_casted->Dtype()}, + [=]() { + grad_input->Fill(0.0); + const Tinput *output_grad_ptr = static_cast(grad_output->DataPtr()); + const Ttarget *target_ptr = static_cast(target->DataPtr()); + const Tinput *input_ptr = static_cast(input_casted->DataPtr()); + Tinput *input_grad_ptr = static_cast(grad_input->DataPtr()); + + CrossEntropyBackwardKernel + <<>>(input_ptr, input_grad_ptr, target_ptr, + output_grad_ptr, bs, num_classes); + }, + "HIP CrossEntropyBackward"); + + return {grad_input}; +} +} // namespace infini_train::kernels::dcu + +#define REGISTER_HIP_CROSS_ENTROPY_KERNEL(kernel_name) \ + REGISTER_KERNEL(infini_train::Device::DeviceType::kDCU, kernel_name, infini_train::kernels::dcu::kernel_name) + +REGISTER_HIP_CROSS_ENTROPY_KERNEL(CrossEntropyForward) +REGISTER_HIP_CROSS_ENTROPY_KERNEL(CrossEntropyBackward) + +#undef REGISTER_HIP_CROSS_ENTROPY_KERNEL diff --git a/infini_train/src/kernels/dcu/elementwise.hip b/infini_train/src/kernels/dcu/elementwise.hip new file mode 100644 index 00000000..4156581c --- /dev/null +++ b/infini_train/src/kernels/dcu/elementwise.hip @@ -0,0 +1,1264 @@ +#include + +#include + +#include "infini_train/include/common/dcu/common_dcu.h" +#include "infini_train/include/common/dcu/kernel_helper.cuh" +#include "infini_train/include/core/runtime/device_guard.h" +#include "infini_train/include/dispatcher.h" +#include "infini_train/include/dtype_dispatch.h" +#include "infini_train/include/tensor.h" + +#include "infini_train/src/core/runtime/dcu/dcu_runtime_common.h" + +namespace infini_train::kernels::dcu { +namespace { +using namespace infini_train::common::dcu; +constexpr int kWarpSize = warpSize; + +// Aligned vector type for vectorized loads/stores (128-bit). +template struct __align__(sizeof(T) * N) aligned_vector { T val[N]; }; + +// Elements per vectorized load/store: 128-bit / sizeof(T). +// float → 4, bf16/half → 8, double → 2. +template constexpr int kVecSize = 16 / sizeof(T); + +// Maximum number of dimensions supported by the broadcast metadata. +// Real-world tensors in this codebase top out at 4-5 dims, so 8 leaves comfortable headroom +// while keeping the struct under the 4 KB HIP kernel parameter limit. +constexpr int kMaxBroadcastDims = 8; + +// POD metadata for broadcast kernels. Passed by value into __global__ kernels so the data +// lives in HIP kernel parameter memory (constant cache) instead of being uploaded via a +// per-call hipMallocAsync + hipMemcpyAsync into global memory. +struct BroadcastMeta { + int ndim; + int64_t a_strides[kMaxBroadcastDims]; + int64_t b_strides[kMaxBroadcastDims]; + int64_t out_strides[kMaxBroadcastDims]; + int64_t a_shape[kMaxBroadcastDims]; + int64_t b_shape[kMaxBroadcastDims]; +}; + +// Build a BroadcastMeta on the host from input/output dim vectors. Right-aligns a_dims/b_dims +// to out_dims's rank (the broadcasting convention) and computes contiguous strides for each. +inline BroadcastMeta MakeBroadcastMeta(const std::vector &a_dims, const std::vector &b_dims, + const std::vector &out_dims) { + BroadcastMeta m{}; + const int ndim = static_cast(out_dims.size()); + CHECK_LE(ndim, kMaxBroadcastDims) << "Broadcast ndim exceeds kMaxBroadcastDims (" << kMaxBroadcastDims << ")"; + m.ndim = ndim; + + std::vector a_shape(ndim, 1), b_shape(ndim, 1); + std::copy_backward(a_dims.begin(), a_dims.end(), a_shape.end()); + std::copy_backward(b_dims.begin(), b_dims.end(), b_shape.end()); + + auto a_str = ComputeStrides(a_shape); + auto b_str = ComputeStrides(b_shape); + auto out_str = ComputeStrides(out_dims); + + for (int i = 0; i < ndim; ++i) { + m.a_strides[i] = a_str[i]; + m.b_strides[i] = b_str[i]; + m.out_strides[i] = out_str[i]; + m.a_shape[i] = a_shape[i]; + m.b_shape[i] = b_shape[i]; + } + return m; +} + +template +__global__ void UnaryForwardKernel(T *output, Func fn, size_t num_elements, size_t offset, const T *input) { + size_t idx = blockIdx.x * blockDim.x + threadIdx.x + offset; + + if (idx < num_elements) { + output[idx] = fn(input[idx]); + } +} + +// Helper for broadcast indexing +__device__ inline int64_t CalcOffset(int64_t idx, int ndim, const int64_t *strides, const int64_t *shape, + const int64_t *out_strides) { + int64_t offset = 0; + for (int i = 0; i < ndim; ++i) { + int64_t out_index = (idx / out_strides[i]) % shape[i]; + int64_t index = shape[i] == 1 ? 0 : out_index; + offset += index * strides[i]; + } + return offset; +} + +inline bool ShapesEqual(const std::vector &a, const std::vector &b) { + if (a.size() != b.size()) { + return false; + } + for (size_t i = 0; i < a.size(); ++i) { + if (a[i] != b[i]) { + return false; + } + } + return true; +} + +template +__global__ void BinaryForwardKernel(T *output, Func fn, BroadcastMeta meta, const T *a, const T *b, + size_t num_elements) { + size_t idx = blockIdx.x * blockDim.x + threadIdx.x; + if (idx >= num_elements) { + return; + } + + int64_t a_offset = CalcOffset(idx, meta.ndim, meta.a_strides, meta.a_shape, meta.out_strides); + int64_t b_offset = CalcOffset(idx, meta.ndim, meta.b_strides, meta.b_shape, meta.out_strides); + + output[idx] = fn(a[a_offset], b[b_offset]); +} + +// Fast path: no broadcast, contiguous tensors — skip CalcOffset entirely +template +__global__ void BinaryForwardKernelNoBroadcast(T *__restrict__ output, Func fn, const T *__restrict__ a, + const T *__restrict__ b, size_t num_elements) { + const size_t grid_stride = static_cast(gridDim.x) * blockDim.x; + for (size_t idx = static_cast(blockIdx.x) * blockDim.x + threadIdx.x; idx < num_elements; + idx += grid_stride) { + output[idx] = fn(a[idx], b[idx]); + } +} + +// Fast path backward: no broadcast, contiguous — skip CalcOffset entirely +template +__global__ void BinaryBackwardKernelNoBroadcastFast(T *__restrict__ outA, T *__restrict__ outB, FuncA fn_a, FuncB fn_b, + size_t numel, const T *__restrict__ grad_out, + const T *__restrict__ inA, const T *__restrict__ inB) { + const size_t grid_stride = static_cast(gridDim.x) * blockDim.x; + for (size_t idx = static_cast(blockIdx.x) * blockDim.x + threadIdx.x; idx < numel; idx += grid_stride) { + const T a = inA ? inA[idx] : T(0); + const T b = inB ? inB[idx] : T(0); + outA[idx] = Mul(grad_out[idx], fn_a(a, b)); + outB[idx] = Mul(grad_out[idx], fn_b(a, b)); + } +} + +// Vectorized fast path backward: no broadcast, contiguous. +// Each thread processes VecSize elements using 128-bit loads/stores. +template +__global__ void BinaryBackwardKernelNoBroadcastVectorized(T *__restrict__ outA, T *__restrict__ outB, FuncA fn_a, + FuncB fn_b, size_t numel, const T *__restrict__ grad_out, + const T *__restrict__ inA, const T *__restrict__ inB) { + using VecT = aligned_vector; + const size_t num_vecs = numel / VecSize; + const size_t grid_stride = static_cast(gridDim.x) * blockDim.x; + + for (size_t vid = static_cast(blockIdx.x) * blockDim.x + threadIdx.x; vid < num_vecs; vid += grid_stride) { + const size_t base = vid * VecSize; + + // 128-bit vectorized loads + VecT g_vec = *reinterpret_cast(&grad_out[base]); + VecT a_vec, b_vec; + if (inA) { + a_vec = *reinterpret_cast(&inA[base]); + } else { +#pragma unroll + for (int i = 0; i < VecSize; ++i) { a_vec.val[i] = T(0); } + } + if (inB) { + b_vec = *reinterpret_cast(&inB[base]); + } else { +#pragma unroll + for (int i = 0; i < VecSize; ++i) { b_vec.val[i] = T(0); } + } + + // Element-wise computation + VecT outA_vec, outB_vec; +#pragma unroll + for (int i = 0; i < VecSize; ++i) { + outA_vec.val[i] = Mul(g_vec.val[i], fn_a(a_vec.val[i], b_vec.val[i])); + outB_vec.val[i] = Mul(g_vec.val[i], fn_b(a_vec.val[i], b_vec.val[i])); + } + + // 128-bit vectorized stores + *reinterpret_cast(&outA[base]) = outA_vec; + *reinterpret_cast(&outB[base]) = outB_vec; + } + + // Handle tail elements (numel % VecSize != 0) + const size_t tail_start = num_vecs * VecSize; + for (size_t idx = tail_start + static_cast(blockIdx.x) * blockDim.x + threadIdx.x; idx < numel; + idx += grid_stride) { + const T a = inA ? inA[idx] : T(0); + const T b = inB ? inB[idx] : T(0); + outA[idx] = Mul(grad_out[idx], fn_a(a, b)); + outB[idx] = Mul(grad_out[idx], fn_b(a, b)); + } +} + +// Helper to choose optimal block size based on tensor size +inline size_t ChooseBlockSize(size_t num_elements) { + if (num_elements < 1024) { + return 64; + } + if (num_elements < 65536) { + return 128; + } + if (num_elements < 1048576) { + return 256; + } + return 512; +} + +// launch the given kernel function with the given output and inputs +template +void LaunchKernel(Kernel &&kernel, const std::shared_ptr &output, const Inputs &...inputs) { + auto extract_ptrs + = [](const auto &...ts) { return std::make_tuple(static_cast(ts ? ts->DataPtr() : nullptr)...); }; + auto input_ptrs = extract_ptrs(inputs...); + + const size_t num_elements = output->NumElements(); + // Use dynamic block size based on tensor size for better occupancy + size_t block_size = std::min(ChooseBlockSize(num_elements), static_cast(1024)); + dim3 block_dims(block_size); + dim3 grid_dims(CEIL_DIV(num_elements, block_dims.x)); + const size_t step = grid_dims.x * block_dims.x; + + for (size_t offset = 0; offset < num_elements; offset += step) { + std::apply([&](auto... ptrs) { kernel(grid_dims, block_dims, offset, ptrs...); }, input_ptrs); + } +} + +// launch a forward elementwise operation given the calculation function, output, and the inputs +// Note: currently only support unary and binary operations +template +void LaunchForward(Func func, const std::shared_ptr &output, const Inputs &...inputs) { + auto device = output->GetDevice(); + const auto &hip_stream = dynamic_cast( + infini_train::core::GetDeviceGuardImpl(device.type())->GetStream(device)) + ->hip_stream(); + T *output_ptr = static_cast(output->DataPtr()); + + if constexpr (sizeof...(inputs) == 1) { + // Unary case + LaunchKernel( + [&](dim3 grid, dim3 block, size_t offset, auto... ptrs) { + UnaryForwardKernel<<>>(output_ptr, func, output->NumElements(), offset, + ptrs...); + }, + output, inputs...); + } else if constexpr (sizeof...(inputs) == 2) { + // Binary case + auto input_tuple = std::make_tuple(inputs...); + const auto &input_a = std::get<0>(input_tuple); + const auto &input_b = std::get<1>(input_tuple); + + const auto &a_dims = input_a->Dims(); + const auto &b_dims = input_b->Dims(); + const auto &out_dims = output->Dims(); + + // Fast path: no broadcast, contiguous — skip hipMalloc/Memcpy/CalcOffset. + // The IsContiguous() guards ensure non-contiguous tensors fall back to the broadcast + // path, keeping the fast path correct when non-contiguous support is added later. + if (ShapesEqual(a_dims, out_dims) && ShapesEqual(b_dims, out_dims) && input_a->IsContiguous() + && input_b->IsContiguous()) { + const size_t num_elements = output->NumElements(); + const T *a_ptr = static_cast(input_a->DataPtr()); + const T *b_ptr = static_cast(input_b->DataPtr()); + dim3 block_dims(std::min(BLOCK_SIZE, static_cast(1024))); + dim3 grid_dims(std::min(CEIL_DIV(num_elements, block_dims.x), static_cast(65535))); + BinaryForwardKernelNoBroadcast<<>>(output_ptr, func, a_ptr, b_ptr, + num_elements); + } else { + // Broadcast path: pass strides/shapes by value via kernel parameter memory. + // This avoids the per-call hipMallocAsync/hipMemcpyAsync/hipFreeAsync that previously + // dominated the host-side jitter floor (especially under LoRA training). + BroadcastMeta meta = MakeBroadcastMeta(a_dims, b_dims, out_dims); + + LaunchKernel( + [&](dim3 grid, dim3 block, size_t /*offset*/, const T *a_ptr, const T *b_ptr) { + BinaryForwardKernel<<>>(output_ptr, func, meta, a_ptr, b_ptr, + output->NumElements()); + }, + output, inputs...); + } + } else { + static_assert(sizeof...(inputs) == 1 || sizeof...(inputs) == 2, + "LaunchForward currently only supports unary and binary operations."); + } +} + +// Backward kernel for unary operators +template +__global__ void UnaryBackwardKernel(T *output, Func fn, size_t num_elements, size_t offset, const T *grad_output, + const T *input) { + size_t idx = blockIdx.x * blockDim.x + threadIdx.x + offset; + + if (idx < num_elements) { + output[idx] = Mul(grad_output[idx], fn(input ? input[idx] : T(0))); + } +} + +enum class BF16Path { NoBroadcast, TwoPassHist, BlockReduce }; + +// Lightweight and stable selector for bf16/half execution paths. +inline BF16Path DecideBF16Path(const std::vector &b_shape, const std::vector &out_shape, + size_t b_num_elements) { + if (ShapesEqual(b_shape, out_shape)) { + return BF16Path::NoBroadcast; + } + const bool varies_last = (b_shape.back() > 1); + if (varies_last) { + if (b_num_elements <= 4096) { + return BF16Path::TwoPassHist; // shared histogram two-pass path + } + } + return BF16Path::BlockReduce; // fallback to block reduction kernel otherwise +} + +// Each B element is used exactly once, so gradients can be written directly without reduction. +template +__global__ void BinaryBackwardKernelNoBroadcast(T *__restrict__ outA, T *__restrict__ outB, FuncA fn_a, FuncB fn_b, + BroadcastMeta meta, size_t numel, const T *__restrict__ grad_out, + const T *__restrict__ inA, const T *__restrict__ inB) { + const size_t grid_stride = static_cast(gridDim.x) * blockDim.x; + for (size_t idx = static_cast(blockIdx.x) * blockDim.x + threadIdx.x; idx < numel; idx += grid_stride) { + const int64_t a_off = CalcOffset(idx, meta.ndim, meta.a_strides, meta.a_shape, meta.out_strides); + const int64_t b_off = CalcOffset(idx, meta.ndim, meta.b_strides, meta.b_shape, meta.out_strides); + + const T a = inA ? inA[a_off] : T(0); + const T b = inB ? inB[b_off] : T(0); + + // Gradient for A has a one-to-one mapping, so we write directly. + outA[a_off] = Mul(grad_out[idx], fn_a(a, b)); + + // Gradient for B also maps one-to-one; no atomics or reductions are required. + outB[b_off] = common::dcu::Cast(Mul(grad_out[idx], fn_b(a, b))); + } +} + +// First pass of histogram two-pass strategy: per-block accumulation in shared memory. +template +__global__ void BinaryBackwardBhistPass1Kernel(T *__restrict__ outA, float *__restrict__ work, FuncA fn_a, FuncB fn_b, + BroadcastMeta meta, size_t numel, int K, const T *__restrict__ grad_out, + const T *__restrict__ inA, const T *__restrict__ inB) { + extern __shared__ float s_hist[]; // dynamic shared memory: K bins plus padding for every 32 buckets + const int pad = K >> 5; // insert one padding slot for every 32 buckets + const int hist_len = K + pad; + + // Zero the shared histogram buffer. + for (int t = threadIdx.x; t < hist_len; t += blockDim.x) { s_hist[t] = 0.0f; } + __syncthreads(); + + const size_t total_threads = (size_t)gridDim.x * blockDim.x; + for (size_t idx = static_cast(blockIdx.x) * blockDim.x + threadIdx.x; idx < numel; idx += total_threads) { + // Linearized offset for B under general broadcasting. + const int64_t b_off = CalcOffset(idx, meta.ndim, meta.b_strides, meta.b_shape, meta.out_strides); + const int bin = static_cast(b_off); // assume K fits in a 32-bit int + const int pbin = bin + (bin >> 5); // apply padding mapping + + // Compute the offset for A under broadcasting. + const int64_t a_off = CalcOffset(idx, meta.ndim, meta.a_strides, meta.a_shape, meta.out_strides); + + const T a = inA ? inA[a_off] : T(0); + const T b = inB ? inB[bin] : T(0); // B is indexed via the flattened bin + + // A is not broadcast, so gradients can be written directly. + outA[a_off] = Mul(grad_out[idx], fn_a(a, b)); + + // Accumulate B's contribution into the shared histogram using float precision. + const float g = common::dcu::Cast(Mul(grad_out[idx], fn_b(a, b))); + atomicAdd(&s_hist[pbin], g); + } + __syncthreads(); + + // Write this block's histogram back to the global workspace: work[block, :]. + float *dst = work + static_cast(blockIdx.x) * static_cast(K); + for (int bin = threadIdx.x; bin < K; bin += blockDim.x) { + const int pbin = bin + (bin >> 5); + dst[bin] = s_hist[pbin]; + } +} + +// Second pass for histogram path: tile the workspace along CTA dimension and atomically add into float buffer. +template +__global__ void BinaryBackwardBhistPass2Reduce2D(const float *__restrict__ work, float *__restrict__ outB_accum, + size_t numBlocks, int K, int tile_height) { + const int k = blockIdx.x * blockDim.x + threadIdx.x; + if (k >= K) { + return; + } + + const size_t begin_row = static_cast(blockIdx.y) * static_cast(tile_height); + const size_t end_row = min(begin_row + static_cast(tile_height), numBlocks); + + float acc = 0.0f; + for (size_t row = begin_row; row < end_row; ++row) { acc += work[row * static_cast(K) + k]; } + + atomicAdd(outB_accum + k, acc); +} + +// Convert the accumulated float buffer back to the target type (bf16/half/float). +template __global__ void CastFloatToTBhist(const float *__restrict__ src, T *__restrict__ dst, int K) { + const int k = blockIdx.x * blockDim.x + threadIdx.x; + if (k < K) { + dst[k] = common::dcu::Cast(src[k]); + } +} + +// Legacy single-dimensional reduction fallback for small grids where atomic tiling is unnecessary. +template +__global__ void BinaryBackwardBhistPass2Reduce1D(const float *__restrict__ work, T *__restrict__ outB, size_t numBlocks, + int K) { + const size_t k = static_cast(blockIdx.x) * blockDim.x + threadIdx.x; + if (k >= static_cast(K)) { + return; + } + + float acc = 0.0f; + for (size_t b = 0; b < numBlocks; ++b) { acc += work[b * static_cast(K) + k]; } + outB[k] = common::dcu::Cast(acc); +} + +// Helper that materializes the two-pass histogram path for bf16/half B gradients. +template +void BinaryBackwardBhistLaunch(FuncA fn_a, FuncB fn_b, T *outA, T *outB, const T *grad_out, const BroadcastMeta &meta, + size_t numel, int K, const T *inA, const T *inB, hipStream_t stream) { + const int kBlockSize = 256; + int grid = static_cast((numel + kBlockSize - 1) / kBlockSize); + if (grid < 1) { + grid = 1; + } + + // Workspace layout: [grid, K] floats. + float *work = nullptr; + HIP_CHECK(hipMallocAsync(&work, static_cast(grid) * static_cast(K) * sizeof(float), stream)); + + // Pass 1: per-block histogram accumulation. + const size_t smem_bytes = static_cast(K + (K >> 5)) * sizeof(float); + BinaryBackwardBhistPass1Kernel + <<>>(outA, work, fn_a, fn_b, meta, numel, K, grad_out, inA, inB); + HIP_CHECK(hipGetLastError()); + + // Pass 2: choose between 1D and 2D reductions depending on workload shape. + int dev = 0; + int sm_count = 0; + HIP_CHECK(hipGetDevice(&dev)); + HIP_CHECK(hipDeviceGetAttribute(&sm_count, hipDeviceAttributeMultiprocessorCount, dev)); + + const int RED_THREADS = 256; + const int oneD_blocks = (K + RED_THREADS - 1) / RED_THREADS; + + // Use the 2D path when the 1D kernel underutilizes the SMs and there are many partial histograms to merge. + const bool use2D = (oneD_blocks < sm_count) && (grid > 4 * sm_count); + + if (!use2D) { + // Fallback: reuse the legacy 1D kernel without atomics. + const dim3 rgrid(oneD_blocks); + const dim3 rblock(RED_THREADS); + BinaryBackwardBhistPass2Reduce1D<<>>(work, outB, static_cast(grid), K); + HIP_CHECK(hipGetLastError()); + } else { + // 2D tiling path: slice the workspace and accumulate using float atomics. + constexpr int kTileHeight = 128; // rows per CTA; tune between 128 and 256 if needed + float *outB_accum = nullptr; + HIP_CHECK(hipMallocAsync(&outB_accum, static_cast(K) * sizeof(float), stream)); + HIP_CHECK(hipMemsetAsync(outB_accum, 0, static_cast(K) * sizeof(float), stream)); + + const dim3 rblock(RED_THREADS, 1, 1); + const dim3 rgrid2((K + RED_THREADS - 1) / RED_THREADS, (grid + kTileHeight - 1) / kTileHeight, 1); + + BinaryBackwardBhistPass2Reduce2D + <<>>(work, outB_accum, static_cast(grid), K, kTileHeight); + HIP_CHECK(hipGetLastError()); + + // Convert accumulated floats back to the target dtype. + const dim3 cgrid((K + RED_THREADS - 1) / RED_THREADS); + CastFloatToTBhist<<>>(outB_accum, outB, K); + HIP_CHECK(hipGetLastError()); + + HIP_CHECK(hipFree(outB_accum)); + } + + HIP_CHECK(hipFree(work)); +} + +// Backward kernel for binary operators +// TODO(lzm): determining and passing b_is_broadcasted from the caller; optimize further +template +__global__ void BinaryBackwardKernel(T *output_a, T *output_b, FuncA fn_a, FuncB fn_b, BroadcastMeta meta, + size_t num_elements, const T *grad_output, const T *input_a, const T *input_b) { + extern __shared__ char shared_memory[]; + const int tid = threadIdx.x; + const int warp_id = tid / kWarpSize; + const int lane_id = tid % kWarpSize; + + using WarpReduce = hipcub::WarpReduce; + WarpReduce::TempStorage *temp_storage = reinterpret_cast(shared_memory); + + size_t idx = blockIdx.x * blockDim.x + tid; + bool in_bounds = (idx < num_elements); + + int64_t a_offset = 0, b_offset = 0; + T a_val = T(0), b_val = T(0); + float grad_val = 0.0f; + + if (in_bounds) { + a_offset = CalcOffset(idx, meta.ndim, meta.a_strides, meta.a_shape, meta.out_strides); + b_offset = CalcOffset(idx, meta.ndim, meta.b_strides, meta.b_shape, meta.out_strides); + a_val = input_a ? input_a[a_offset] : T(0); + b_val = input_b ? input_b[b_offset] : T(0); + output_a[a_offset] = Mul(grad_output[idx], fn_a(a_val, b_val)); + grad_val = common::dcu::Cast(Mul(grad_output[idx], fn_b(a_val, b_val))); + } + + unsigned long long active_mask = __ballot(in_bounds); + if (!active_mask) { + return; + } + + int leader = __ffsll(active_mask) - 1; + int64_t common_offset = __shfl(b_offset, leader); + + // Check if all active threads share common b_offset + bool warp_uniform = true; + for (int i = 0; i < kWarpSize; ++i) { + if (!(active_mask & (1ULL << i))) { + continue; + } + int64_t offset_i = __shfl(b_offset, i); + if (offset_i != common_offset) { + warp_uniform = false; + break; + } + } + + if (warp_uniform) { + float reduced = WarpReduce(temp_storage[warp_id]).Sum(grad_val); + if (lane_id == leader) { + // FIXME(lzm): atomicAdd is much slower for bf16 and half compared to float, needs further optimization + common::dcu::AtomicAdd(&output_b[common_offset], common::dcu::Cast(reduced)); + } + } else if (in_bounds) { + // FIXME(lzm): atomicAdd is much slower for bf16 and half compared to float, needs further optimization + common::dcu::AtomicAdd(&output_b[b_offset], common::dcu::Cast(grad_val)); + } +} + +// NOTE(dcj): Specialized BinaryBackwardKernel for low-precision types (__half / bfloat16) +template +__global__ void BinaryBackwardKernel(T *output_a, T *output_b, FuncA fn_a, FuncB fn_b, BroadcastMeta meta, + size_t num_elements, size_t b_num_elements, const T *grad_output, const T *input_a, + const T *input_b, bool fast_atomics) { + + const int tid = threadIdx.x; + const int block_threads = blockDim.x; + const int global_idx = blockIdx.x * blockDim.x + tid; + bool in_bounds = (global_idx < num_elements); + + // Dynamic shared memory layout: split offsets and gradients into parallel arrays. + extern __shared__ char shared_memory[]; + int64_t *s_offset = reinterpret_cast(shared_memory); + float *s_grad = reinterpret_cast(s_offset + block_threads + block_threads / kWarpSize); + + // Padding: insert one slot per 32 threads to avoid bank conflicts. + const int padded_tid = tid + (tid >> 5); + + // Each thread calculates its own a_offset and b_offset + int64_t a_offset = 0, b_offset = 0; + float grad_val = 0.0f; + T a_val = T(0), b_val = T(0); + + if (in_bounds) { + a_offset = CalcOffset(global_idx, meta.ndim, meta.a_strides, meta.a_shape, meta.out_strides); + b_offset = CalcOffset(global_idx, meta.ndim, meta.b_strides, meta.b_shape, meta.out_strides); + + a_val = input_a ? input_a[a_offset] : T(0); + b_val = input_b ? input_b[b_offset] : T(0); + + // Compute gradient contribution for output_a + output_a[a_offset] = Mul(grad_output[global_idx], fn_a(a_val, b_val)); + // Store gradient contribution for output_b in float for accumulation + grad_val = common::dcu::Cast(Mul(grad_output[global_idx], fn_b(a_val, b_val))); + } + + // Store partial results in shared memory. + s_offset[padded_tid] = in_bounds ? b_offset : -1; + s_grad[padded_tid] = grad_val; + + __syncthreads(); + + // Perform block-wide reduction with padded indices. + for (int stride = 1; stride < block_threads; stride *= 2) { + __syncthreads(); + if ((tid % (2 * stride)) == 0 && (tid + stride) < block_threads) { + const int p1 = tid + (tid >> 5); + const int p2 = (tid + stride) + ((tid + stride) >> 5); + + if (s_offset[p1] == s_offset[p2] && s_offset[p1] != -1) { + s_grad[p1] += s_grad[p2]; + s_offset[p2] = -1; + } + } + } + __syncthreads(); + + // Write final result back to global memory + if (in_bounds) { + const int shared_idx = tid + (tid >> 5); + if (s_offset[shared_idx] != -1) { + fastAtomicAdd(output_b, s_offset[shared_idx], b_num_elements, + common::dcu::Cast(s_grad[shared_idx]), fast_atomics); + } + } +} + +// launch unary operator's backward kernel +template +void LaunchBackward(Func func, const std::shared_ptr &output, const std::shared_ptr &grad_output, + const Inputs &...inputs) { + auto device = output->GetDevice(); + const auto &hip_stream = dynamic_cast( + infini_train::core::GetDeviceGuardImpl(device.type())->GetStream(device)) + ->hip_stream(); + + T *output_ptr = static_cast(output->DataPtr()); + const T *grad_ptr = static_cast(grad_output->DataPtr()); + + LaunchKernel( + [=](dim3 grid, dim3 block, size_t offset, auto... ptrs) { + UnaryBackwardKernel<<>>(output_ptr, func, output->NumElements(), offset, + grad_ptr, ptrs...); + }, + output, inputs...); +} + +// launch binary operator's backward kernel +template +void LaunchBackward(FuncA fun_a, FuncB fun_b, const std::shared_ptr &output_a, + const std::shared_ptr &output_b, const std::vector &a_dims, + const std::vector &b_dims, const std::shared_ptr &grad_output, + const Inputs &...inputs) { + auto device = output_a->GetDevice(); + const auto &stream = dynamic_cast( + infini_train::core::GetDeviceGuardImpl(device.type())->GetStream(device)) + ->hip_stream(); + + T *output_a_ptr = static_cast(output_a->DataPtr()); + T *output_b_ptr = static_cast(output_b->DataPtr()); + const T *grad_output_ptr = static_cast(grad_output->DataPtr()); + + const auto &out_dims = grad_output->Dims(); + const size_t num_elements = grad_output->NumElements(); + + // Fast path: no broadcast, contiguous — skip hipMalloc/Memcpy/CalcOffset. + // The IsContiguous() guard ensures non-contiguous grad_output falls back to the broadcast + // path, keeping the fast path correct when non-contiguous support is added later. + if (ShapesEqual(a_dims, b_dims) && ShapesEqual(a_dims, out_dims) && grad_output->IsContiguous()) { + auto extract_ptrs = [](const auto &...ts) { + return std::make_tuple(static_cast(ts ? ts->DataPtr() : nullptr)...); + }; + auto [input_a_ptr, input_b_ptr] = extract_ptrs(inputs...); + + constexpr int VecSize = kVecSize; + // Use vectorized kernel if all pointers are 16-byte aligned and numel is large enough + const bool can_vectorize + = (num_elements >= static_cast(VecSize)) + && (reinterpret_cast(output_a_ptr) % (sizeof(T) * VecSize) == 0) + && (reinterpret_cast(output_b_ptr) % (sizeof(T) * VecSize) == 0) + && (reinterpret_cast(grad_output_ptr) % (sizeof(T) * VecSize) == 0) + && (!input_a_ptr || reinterpret_cast(input_a_ptr) % (sizeof(T) * VecSize) == 0) + && (!input_b_ptr || reinterpret_cast(input_b_ptr) % (sizeof(T) * VecSize) == 0); + + if (can_vectorize) { + const size_t num_vecs = num_elements / VecSize; + dim3 block_dims(std::min(static_cast(256), std::min(num_vecs, static_cast(1024)))); + dim3 grid_dims(std::min(CEIL_DIV(num_vecs, block_dims.x), static_cast(65535))); + BinaryBackwardKernelNoBroadcastVectorized<<>>( + output_a_ptr, output_b_ptr, fun_a, fun_b, num_elements, grad_output_ptr, input_a_ptr, input_b_ptr); + } else { + dim3 block_dims(std::min(BLOCK_SIZE, static_cast(1024))); + dim3 grid_dims(std::min(CEIL_DIV(num_elements, block_dims.x), static_cast(65535))); + BinaryBackwardKernelNoBroadcastFast<<>>( + output_a_ptr, output_b_ptr, fun_a, fun_b, num_elements, grad_output_ptr, input_a_ptr, input_b_ptr); + } + return; + } + + // Broadcast path: pass strides/shapes by value via kernel parameter memory. + // This avoids the per-call hipMallocAsync/hipMemcpyAsync/hipFreeAsync that previously + // dominated the host-side jitter floor (especially under LoRA training). + BroadcastMeta meta = MakeBroadcastMeta(a_dims, b_dims, out_dims); + + if constexpr (std::is_same_v) { + LaunchKernel( + [=](dim3 grid, dim3 block, size_t /*offset*/, auto... ptrs) { + const int num_warps = BLOCK_SIZE / kWarpSize; + const size_t smem_size = num_warps * sizeof(hipcub::WarpReduce::TempStorage); + BinaryBackwardKernel<<>>(output_a_ptr, output_b_ptr, fun_a, fun_b, meta, + num_elements, grad_output_ptr, ptrs...); + }, + output_a, inputs...); + } else if constexpr (std::is_same_v || std::is_same_v) { + // Dynamically choose the most efficient bf16/half strategy based on broadcast pattern. + // Reconstruct right-aligned b_shape (stack-only, no device allocations) for + // DecideBF16Path which still operates on std::vector. + const int ndim = meta.ndim; + std::vector b_shape(meta.b_shape, meta.b_shape + ndim); + const std::vector &out_shape = out_dims; + + size_t b_num_elements = 1; + for (auto v : b_shape) { b_num_elements *= static_cast(v); } + const int K_linear = static_cast(b_num_elements); + + // Select the execution path. + const BF16Path path = DecideBF16Path(b_shape, out_shape, b_num_elements); + + if (path == BF16Path::NoBroadcast) { + // No broadcast: write gradients directly without shared memory or atomics. + LaunchKernel( + [=](dim3 grid, dim3 block, size_t /*offset*/, auto... ptrs) { + BinaryBackwardKernelNoBroadcast<<>>( + output_a_ptr, output_b_ptr, fun_a, fun_b, meta, num_elements, grad_output_ptr, ptrs...); + }, + output_a, inputs...); + return; + } + + if (path == BF16Path::TwoPassHist) { + // Small K with variation in the innermost dimension: use two-pass histogram strategy. + LaunchKernel( + [=](dim3 /*grid*/, dim3 /*block*/, size_t /*offset*/, const T *input_a_ptr, const T *input_b_ptr) { + BinaryBackwardBhistLaunch(fun_a, fun_b, output_a_ptr, output_b_ptr, + grad_output_ptr, meta, num_elements, K_linear, + input_a_ptr, input_b_ptr, stream); + }, + output_a, inputs...); + + return; + } + + // Otherwise fall back to the block-reduction kernel with SoA layout and fast atomics. + LaunchKernel( + [=](dim3 grid, dim3 block, size_t /*offset*/, auto... ptrs) { + const int padded_block = BLOCK_SIZE + BLOCK_SIZE / kWarpSize; + const size_t smem_size = static_cast(padded_block) * (sizeof(int64_t) + sizeof(float)); + BinaryBackwardKernel<<>>( + output_a_ptr, output_b_ptr, fun_a, fun_b, meta, num_elements, output_b->NumElements(), + grad_output_ptr, ptrs..., /*fast_atomics=*/true); + }, + output_a, inputs...); + } +} + +template std::shared_ptr UnaryForward(const std::shared_ptr &input, Func unary_fn) { + auto dtype = input->Dtype(); + auto output = std::make_shared(input->Dims(), dtype, input->GetDevice()); + + switch (dtype) { + DISPATCH_CASE(WRAP(LaunchForward<256, float>(unary_fn, output, input);), DataType::kFLOAT32) + DISPATCH_CASE(WRAP(LaunchForward<256, hip_bfloat16>(unary_fn, output, input);), DataType::kBFLOAT16) + DISPATCH_CASE(WRAP(LaunchForward<256, int64_t>(unary_fn, output, input);), DataType::kINT64) + default: + LOG_LOC(FATAL, "HIP unary forward: 'Unsupported data type'"); + } + + return output; +} + +template +std::shared_ptr UnaryBackward(const std::shared_ptr &grad_output, const std::shared_ptr &a, + Func unary_fn) { + auto dtype = grad_output->Dtype(); + auto a_dtype = a ? a->Dtype() : dtype; + DataType promoted_type = PromoteDataTypes(dtype, a_dtype); + + auto grad_output_promoted + = dtype == promoted_type ? grad_output : std::make_shared(grad_output->To(promoted_type)); + auto a_promoted = a_dtype == promoted_type ? a : std::make_shared(a->To(promoted_type)); + auto output = std::make_shared(grad_output->Dims(), promoted_type, grad_output->GetDevice()); + + switch (promoted_type) { + DISPATCH_CASE(WRAP({ LaunchBackward<256, float>(unary_fn, output, grad_output_promoted, a_promoted); }), + DataType::kFLOAT32) + DISPATCH_CASE(WRAP({ LaunchBackward<256, hip_bfloat16>(unary_fn, output, grad_output_promoted, a_promoted); }), + DataType::kBFLOAT16) + DISPATCH_CASE(WRAP({ LaunchBackward<256, int64_t>(unary_fn, output, grad_output_promoted, a_promoted); }), + DataType::kINT64) + default: + LOG_LOC(FATAL, "HIP unary backward: 'Unsupported data type'"); + } + + return output; +} + +template +std::shared_ptr BinaryForward(const std::shared_ptr &a, const std::shared_ptr &b, + Func binary_fn) { + auto a_dtype = a->Dtype(); + auto b_dtype = b->Dtype(); + + DataType promoted_type = PromoteDataTypes(a_dtype, b_dtype); + + auto a_promoted = a_dtype == promoted_type ? a : std::make_shared(a->To(promoted_type)); + auto b_promoted = b_dtype == promoted_type ? b : std::make_shared(b->To(promoted_type)); + // Currently a and b should have the same data type and only one-way broadcasting from b to a is assumed by + // default + CHECK(a->NumElements() >= b->NumElements() && a->NumElements() % b->NumElements() == 0); + + auto output = std::make_shared(a->Dims(), promoted_type, a->GetDevice()); + + switch (promoted_type) { + DISPATCH_CASE(WRAP(LaunchForward<256, float>(binary_fn, output, a_promoted, b_promoted);), DataType::kFLOAT32) + DISPATCH_CASE(WRAP(LaunchForward<256, hip_bfloat16>(binary_fn, output, a_promoted, b_promoted);), + DataType::kBFLOAT16) + DISPATCH_CASE(WRAP(LaunchForward<256, int64_t>(binary_fn, output, a_promoted, b_promoted);), DataType::kINT64) + default: + LOG_LOC(FATAL, "HIP binary forward: 'Unsupported data type'"); + } + + return output; +} + +template +std::pair, std::shared_ptr> +BinaryBackward(const std::shared_ptr &grad_output, const std::shared_ptr &a, + const std::shared_ptr &b, const std::vector &a_dims, const std::vector &b_dims, + FuncA fn_a, FuncB fn_b) { + const auto a_num_elements = std::accumulate(a_dims.begin(), a_dims.end(), 1, std::multiplies()); + const auto b_num_elements = std::accumulate(b_dims.begin(), b_dims.end(), 1, std::multiplies()); + + std::shared_ptr a_promoted = a; + std::shared_ptr b_promoted = b; + std::shared_ptr grad_output_promoted = grad_output; + + auto dtype = grad_output_promoted->Dtype(); + auto device = grad_output->GetDevice(); + + auto a_dtype = a_promoted ? a_promoted->Dtype() : dtype; + auto b_dtype = b_promoted ? b_promoted->Dtype() : dtype; + // Compute dtype determined by saved tensors (forward compute dtype), not grad_output + DataType promoted_type = PromoteDataTypes(a_dtype, b_dtype); + + CHECK(a_num_elements >= b_num_elements && a_num_elements % b_num_elements == 0); + + auto promote_if_needed = [&](std::shared_ptr &t, size_t expected_numel, DataType promoted_type) { + if (t) { + CHECK(expected_numel == t->NumElements()); + if (t->Dtype() != promoted_type) { + t = std::make_shared(t->To(promoted_type)); + } + } + }; + promote_if_needed(a_promoted, a_num_elements, promoted_type); + promote_if_needed(b_promoted, b_num_elements, promoted_type); + if (dtype != promoted_type) { + grad_output_promoted = std::make_shared(grad_output_promoted->To(promoted_type)); + } + + auto grad_a = std::make_shared(a_dims, promoted_type, device); + auto grad_b = std::make_shared(b_dims, promoted_type, device); + + // Only Fill(0) when broadcast is needed (atomicAdd requires zero-init). + // The no-broadcast fast path writes every element directly. + const bool needs_broadcast = !ShapesEqual(a_dims, b_dims) || !ShapesEqual(a_dims, grad_output->Dims()); + + switch (promoted_type) { + DISPATCH_CASE(WRAP({ + if (needs_broadcast) { + grad_a->Fill(0.0f); + grad_b->Fill(0.0f); + } + LaunchBackward<256, float>(fn_a, fn_b, grad_a, grad_b, a_dims, b_dims, grad_output_promoted, + a_promoted, b_promoted); + }), + DataType::kFLOAT32) + DISPATCH_CASE(WRAP({ + if (needs_broadcast) { + grad_a->Fill(0.0f); + grad_b->Fill(0.0f); + } + LaunchBackward<256, hip_bfloat16>(fn_a, fn_b, grad_a, grad_b, a_dims, b_dims, + grad_output_promoted, a_promoted, b_promoted); + }), + DataType::kBFLOAT16) + // FIXME(zbl): AtomicAdd does not support int64_t + // DISPATCH_CASE(WRAP({ + // grad_a->Fill(0.0); + // grad_b->Fill(0.0); + // LaunchBackward<256, int64_t>(fn_a, fn_b, grad_a, grad_b, a_dims, b_dims, grad_output, a, + // b); + // }), + // DataType::kINT64) + default: + LOG_LOC(FATAL, "HIP binary backward: 'Unsupported data type'"); + } + + return {grad_a, grad_b}; +} +} // namespace + +std::shared_ptr NegForward(const std::shared_ptr &input) { + DISPATCH(input->Dtype(), return UnaryForward(input, [] __device__(auto x) { return Neg(x); }); + , INFINI_ALL_FLOATING_TYPES) +} + +std::shared_ptr NegBackward(const std::shared_ptr &grad_output) { + DISPATCH(grad_output->Dtype(), + return UnaryBackward(grad_output, nullptr, [] __device__(auto x) { return decltype(x){-1}; }); + , INFINI_ALL_FLOATING_TYPES) +} + +std::shared_ptr ReciprocalForward(const std::shared_ptr &input) { + DISPATCH(input->Dtype(), return UnaryForward(input, [] __device__(auto x) { return Reciprocal(x); }); + , INFINI_ALL_FLOATING_TYPES) +} + +std::shared_ptr ReciprocalBackward(const std::shared_ptr &grad_output, + const std::shared_ptr &input) { + DISPATCH( + grad_output->Dtype(), + return UnaryBackward(grad_output, input, [] __device__(auto x) { return Div(decltype(x){-1}, Mul(x, x)); }); + , INFINI_ALL_FLOATING_TYPES) +} + +std::shared_ptr SinForward(const std::shared_ptr &input) { + DISPATCH(input->Dtype(), return UnaryForward(input, [] __device__(auto x) { return Sin(x); }); + , INFINI_ALL_FLOATING_TYPES) +} + +std::shared_ptr SinBackward(const std::shared_ptr &grad_output, const std::shared_ptr &input) { + DISPATCH(grad_output->Dtype(), return UnaryBackward(grad_output, input, [] __device__(auto x) { return Cos(x); }); + , INFINI_ALL_FLOATING_TYPES) +} + +std::shared_ptr CosForward(const std::shared_ptr &input) { + DISPATCH(input->Dtype(), return UnaryForward(input, [] __device__(auto x) { return Cos(x); }); + , INFINI_ALL_FLOATING_TYPES) +} + +std::shared_ptr CosBackward(const std::shared_ptr &grad_output, const std::shared_ptr &input) { + DISPATCH(grad_output->Dtype(), + return UnaryBackward(grad_output, input, [] __device__(auto x) { return Neg(Sin(x)); }); + , INFINI_ALL_FLOATING_TYPES) +} + +std::shared_ptr TanhForward(const std::shared_ptr &input) { + DISPATCH(input->Dtype(), return UnaryForward(input, [] __device__(auto x) { return Tanh(x); }); + , INFINI_ALL_FLOATING_TYPES) +} + +std::shared_ptr TanhBackward(const std::shared_ptr &grad_output, + const std::shared_ptr &output) { + DISPATCH(grad_output->Dtype(), + return UnaryBackward(grad_output, output, [] __device__(auto x) { return decltype(x){1} - Mul(x, x); }); + , INFINI_ALL_FLOATING_TYPES) +} + +std::shared_ptr PowForward(const std::shared_ptr &input, float scalar, bool scalar_is_base) { + DISPATCH(input->Dtype(), WRAP({ + if (scalar_is_base) { + return UnaryForward( + input, [scalar] __device__(auto x) { return Pow(static_cast(scalar), x); }); + } else { + return UnaryForward( + input, [scalar] __device__(auto x) { return Pow(x, static_cast(scalar)); }); + } + }), + INFINI_ALL_FLOATING_TYPES); +} + +std::shared_ptr PowBackward(const std::shared_ptr &grad_output, const std::shared_ptr &input, + float scalar, bool scalar_is_base) { + DISPATCH(grad_output->Dtype(), + return UnaryBackward(grad_output, input, + [scalar, scalar_is_base] __device__(auto x) { + auto casted_scalar = common::dcu::Cast(scalar); + if (scalar_is_base) { + return Mul(Log(casted_scalar), Pow(casted_scalar, x)); + } else { + return Mul(casted_scalar, Pow(x, casted_scalar - decltype(x){1})); + } + }); + , INFINI_ALL_FLOATING_TYPES) +} + +std::shared_ptr RsqrtForward(const std::shared_ptr &input) { + DISPATCH(input->Dtype(), return UnaryForward(input, [] __device__(auto x) { return Rsqrt(x); }); + , INFINI_ALL_FLOATING_TYPES) +} + +std::shared_ptr RsqrtBackward(const std::shared_ptr &grad_output, + const std::shared_ptr &input) { + DISPATCH(grad_output->Dtype(), + return UnaryBackward( + grad_output, input, + [] __device__(auto x) { return Mul(static_cast(-0.5), Mul(Reciprocal(x), Rsqrt(x))); }); + , INFINI_ALL_FLOATING_TYPES) +} + +std::shared_ptr ExpForward(const std::shared_ptr &input) { + DISPATCH(input->Dtype(), return UnaryForward(input, [] __device__(auto x) { return Exp(x); }); + , INFINI_ALL_FLOATING_TYPES) +} + +std::shared_ptr ExpBackward(const std::shared_ptr &grad_output, const std::shared_ptr &output) { + DISPATCH(grad_output->Dtype(), return UnaryBackward(grad_output, output, [] __device__(auto y) { return y; }); + , INFINI_ALL_FLOATING_TYPES) +} + +std::shared_ptr LogForward(const std::shared_ptr &input) { + DISPATCH(input->Dtype(), return UnaryForward(input, [] __device__(auto x) { return Log(x); }); + , INFINI_ALL_FLOATING_TYPES) +} + +std::shared_ptr LogBackward(const std::shared_ptr &grad_output, const std::shared_ptr &input) { + DISPATCH(grad_output->Dtype(), + return UnaryBackward(grad_output, input, [] __device__(auto x) { return Reciprocal(x); }); + , INFINI_ALL_FLOATING_TYPES) +} + +std::shared_ptr EqualsForward(const std::shared_ptr &a, const std::shared_ptr &b) { + DISPATCH(a->Dtype(), + return BinaryForward(a, b, + [] __device__(auto x, auto y) { return (x == y) ? decltype(x){1} : decltype(x){0}; }); + , INFINI_ALL_TYPES) +} + +std::shared_ptr EqualsScalarForward(const std::shared_ptr &a, float scalar) { + DISPATCH(a->Dtype(), return UnaryForward(a, + [scalar] __device__(auto x) { + return x == static_cast(scalar) ? decltype(x){1} + : decltype(x){0}; + }); + , INFINI_ALL_FLOATING_TYPES) +} + +std::shared_ptr LtForward(const std::shared_ptr &a, const std::shared_ptr &b) { + DISPATCH(a->Dtype(), return BinaryForward( + a, b, [] __device__(auto x, auto y) { return x < y ? decltype(x){1} : decltype(x){0}; }); + , INFINI_ALL_TYPES) +} + +std::shared_ptr LtScalarForward(const std::shared_ptr &a, float scalar) { + DISPATCH(a->Dtype(), return UnaryForward(a, + [scalar] __device__(auto x) { + return (x < static_cast(scalar)) ? decltype(x){1} + : decltype(x){0}; + }); + , INFINI_ALL_TYPES) +} + +std::shared_ptr LeForward(const std::shared_ptr &a, const std::shared_ptr &b) { + DISPATCH(a->Dtype(), + return BinaryForward(a, b, + [] __device__(auto x, auto y) { return (x <= y) ? decltype(x){1} : decltype(x){0}; }); + , INFINI_ALL_TYPES) +} + +std::shared_ptr LeScalarForward(const std::shared_ptr &a, float scalar) { + DISPATCH(a->Dtype(), return UnaryForward(a, + [scalar] __device__(auto x) { + return (x <= static_cast(scalar)) ? decltype(x){1} + : decltype(x){0}; + }); + , INFINI_ALL_TYPES) +} + +std::shared_ptr GtForward(const std::shared_ptr &a, const std::shared_ptr &b) { + DISPATCH(a->Dtype(), return BinaryForward( + a, b, [] __device__(auto x, auto y) { return x > y ? decltype(x){1} : decltype(x){0}; }); + , INFINI_ALL_TYPES) +} + +std::shared_ptr GtScalarForward(const std::shared_ptr &a, float scalar) { + DISPATCH(a->Dtype(), return UnaryForward(a, + [scalar] __device__(auto x) { + return (x > static_cast(scalar)) ? decltype(x){1} + : decltype(x){0}; + }); + , INFINI_ALL_TYPES) +} + +std::shared_ptr GeForward(const std::shared_ptr &a, const std::shared_ptr &b) { + DISPATCH(a->Dtype(), + return BinaryForward(a, b, + [] __device__(auto x, auto y) { return (x >= y) ? decltype(x){1} : decltype(x){0}; }); + , INFINI_ALL_TYPES) +} + +std::shared_ptr GeScalarForward(const std::shared_ptr &a, float scalar) { + DISPATCH(a->Dtype(), return UnaryForward(a, + [scalar] __device__(auto x) { + return (x >= static_cast(scalar)) ? decltype(x){1} + : decltype(x){0}; + }); + , INFINI_ALL_TYPES) +} + +std::shared_ptr OrForward(const std::shared_ptr &a, const std::shared_ptr &b) { + DISPATCH(a->Dtype(), return BinaryForward(a, b, + [] __device__(auto x, auto y) { + return (x != decltype(x){0} || y != decltype(y){0}) ? decltype(x){1} + : decltype(x){0}; + }); + , INFINI_ALL_TYPES) +} + +std::shared_ptr AndForward(const std::shared_ptr &a, const std::shared_ptr &b) { + DISPATCH(a->Dtype(), return BinaryForward(a, b, + [] __device__(auto x, auto y) { + return (x != decltype(x){0} && y != decltype(y){0}) ? decltype(x){1} + : decltype(x){0}; + }); + , INFINI_ALL_TYPES) +} + +std::shared_ptr AddForward(const std::shared_ptr &a, const std::shared_ptr &b) { + DISPATCH(a->Dtype(), return BinaryForward(a, b, [] __device__(auto x, auto y) { return Add(x, y); }); + , INFINI_ALL_FLOATING_TYPES) +} + +std::pair, std::shared_ptr> AddBackward(const std::shared_ptr &grad_output, + const std::vector &a_dims, + const std::vector &b_dims) { + auto fn = [] __device__(auto x, auto y) { return decltype(x){1}; }; + return BinaryBackward(grad_output, nullptr, nullptr, a_dims, b_dims, fn, fn); +} + +std::shared_ptr AddScalarForward(const std::shared_ptr &a, float scalar) { + DISPATCH(a->Dtype(), + return UnaryForward(a, [scalar] __device__(auto x) { return Add(x, static_cast(scalar)); }); + , INFINI_ALL_TYPES) +} + +std::shared_ptr AddScalarBackward(const std::shared_ptr &grad_output) { + DISPATCH(grad_output->Dtype(), + return UnaryBackward(grad_output, nullptr, + [] __device__(auto x) { return common::dcu::Cast(1); }); + , INFINI_ALL_TYPES) +} + +std::shared_ptr SubForward(const std::shared_ptr &a, const std::shared_ptr &b) { + DISPATCH(a->Dtype(), return BinaryForward(a, b, [] __device__(auto x, auto y) { return Sub(x, y); }); + , INFINI_ALL_TYPES) +} + +std::pair, std::shared_ptr> SubBackward(const std::shared_ptr &grad_output, + const std::vector &a_dims, + const std::vector &b_dims) { + auto fn_a = [] __device__(auto x, auto y) { return decltype(x){1}; }; + auto fn_b = [] __device__(auto x, auto y) { return decltype(x){-1}; }; + return BinaryBackward(grad_output, nullptr, nullptr, a_dims, b_dims, fn_a, fn_b); +} + +std::shared_ptr MulForward(const std::shared_ptr &a, const std::shared_ptr &b) { + DISPATCH(a->Dtype(), return BinaryForward(a, b, [] __device__(auto x, auto y) { return Mul(x, y); }); + , INFINI_ALL_FLOATING_TYPES) +} + +std::pair, std::shared_ptr> MulBackward(const std::shared_ptr &grad_output, + const std::shared_ptr &a, + const std::shared_ptr &b) { + DISPATCH_WITH_DEFAULT(grad_output->Dtype(), + return BinaryBackward( + grad_output, a, b, a->Dims(), b->Dims(), [] __device__(auto, auto y) { return y; }, + [] __device__(auto x, auto) { return x; }); + , WRAP({ + LOG_LOC(FATAL, "HIP MulBackward: 'Unsupported data type'"); + return {nullptr, nullptr}; + }), + INFINI_ALL_FLOATING_TYPES) +} + +std::shared_ptr MulScalarForward(const std::shared_ptr &a, float scalar) { + DISPATCH(a->Dtype(), + return UnaryForward(a, [scalar] __device__(auto x) { return Mul(x, static_cast(scalar)); }); + , INFINI_ALL_FLOATING_TYPES) +} + +std::shared_ptr MulScalarBackward(const std::shared_ptr &grad_output, float scalar) { + DISPATCH(grad_output->Dtype(), + return UnaryBackward(grad_output, nullptr, + [scalar] __device__(auto x) { return static_cast(scalar); }); + , INFINI_ALL_FLOATING_TYPES) +} + +std::shared_ptr DivForward(const std::shared_ptr &a, const std::shared_ptr &b) { + DISPATCH(a->Dtype(), return BinaryForward(a, b, [] __device__(auto x, auto y) { return Div(x, y); }); + , INFINI_ALL_FLOATING_TYPES) +} + +std::pair, std::shared_ptr> DivBackward(const std::shared_ptr &grad_output, + const std::shared_ptr &a, + const std::shared_ptr &b) { + DISPATCH_WITH_DEFAULT(grad_output->Dtype(), return BinaryBackward( + grad_output, a, b, a->Dims(), b->Dims(), + [] __device__(auto, auto y) { return Reciprocal(y); }, + [] __device__(auto x, auto y) { return Div(Neg(x), Mul(y, y)); }); + , WRAP({ + LOG_LOC(FATAL, "HIP DivBackward: 'Unsupported data type'"); + return {nullptr, nullptr}; + }), + INFINI_ALL_FLOATING_TYPES) +} + +std::shared_ptr SigmoidForward(const std::shared_ptr &input) { + DISPATCH(input->Dtype(), return UnaryForward(input, [] __device__(auto x) { return Sigmoid(x); }); + , INFINI_ALL_FLOATING_TYPES) +} + +std::shared_ptr SigmoidBackward(const std::shared_ptr &output, + const std::shared_ptr &grad_output) { + DISPATCH( + grad_output->Dtype(), + return UnaryBackward(grad_output, output, [] __device__(auto x) { return Mul(x, Sub(decltype(x){1}, x)); }); + , INFINI_ALL_FLOATING_TYPES) +} +} // namespace infini_train::kernels::dcu + +#define REGISTER_HIP_ELEMENTWISE_KERNEL(kernel_name) \ + REGISTER_KERNEL(infini_train::Device::DeviceType::kDCU, kernel_name, infini_train::kernels::dcu::kernel_name) + +REGISTER_HIP_ELEMENTWISE_KERNEL(NegForward) +REGISTER_HIP_ELEMENTWISE_KERNEL(NegBackward) +REGISTER_HIP_ELEMENTWISE_KERNEL(ReciprocalForward) +REGISTER_HIP_ELEMENTWISE_KERNEL(ReciprocalBackward) +REGISTER_HIP_ELEMENTWISE_KERNEL(SinForward) +REGISTER_HIP_ELEMENTWISE_KERNEL(SinBackward) +REGISTER_HIP_ELEMENTWISE_KERNEL(CosForward) +REGISTER_HIP_ELEMENTWISE_KERNEL(CosBackward) +REGISTER_HIP_ELEMENTWISE_KERNEL(TanhForward) +REGISTER_HIP_ELEMENTWISE_KERNEL(TanhBackward) +REGISTER_HIP_ELEMENTWISE_KERNEL(PowForward) +REGISTER_HIP_ELEMENTWISE_KERNEL(PowBackward) +REGISTER_HIP_ELEMENTWISE_KERNEL(RsqrtForward) +REGISTER_HIP_ELEMENTWISE_KERNEL(RsqrtBackward) +REGISTER_HIP_ELEMENTWISE_KERNEL(ExpForward) +REGISTER_HIP_ELEMENTWISE_KERNEL(ExpBackward) +REGISTER_HIP_ELEMENTWISE_KERNEL(LogForward) +REGISTER_HIP_ELEMENTWISE_KERNEL(LogBackward) +REGISTER_HIP_ELEMENTWISE_KERNEL(EqualsForward) +REGISTER_HIP_ELEMENTWISE_KERNEL(EqualsScalarForward) +REGISTER_HIP_ELEMENTWISE_KERNEL(LtForward) +REGISTER_HIP_ELEMENTWISE_KERNEL(LtScalarForward) +REGISTER_HIP_ELEMENTWISE_KERNEL(LeForward) +REGISTER_HIP_ELEMENTWISE_KERNEL(LeScalarForward) +REGISTER_HIP_ELEMENTWISE_KERNEL(GtForward) +REGISTER_HIP_ELEMENTWISE_KERNEL(GtScalarForward) +REGISTER_HIP_ELEMENTWISE_KERNEL(GeForward) +REGISTER_HIP_ELEMENTWISE_KERNEL(GeScalarForward) +REGISTER_HIP_ELEMENTWISE_KERNEL(OrForward) +REGISTER_HIP_ELEMENTWISE_KERNEL(AndForward) +REGISTER_HIP_ELEMENTWISE_KERNEL(AddForward) +REGISTER_HIP_ELEMENTWISE_KERNEL(AddBackward) +REGISTER_HIP_ELEMENTWISE_KERNEL(AddScalarForward) +REGISTER_HIP_ELEMENTWISE_KERNEL(AddScalarBackward) +REGISTER_HIP_ELEMENTWISE_KERNEL(SubForward) +REGISTER_HIP_ELEMENTWISE_KERNEL(SubBackward) +REGISTER_HIP_ELEMENTWISE_KERNEL(MulForward) +REGISTER_HIP_ELEMENTWISE_KERNEL(MulBackward) +REGISTER_HIP_ELEMENTWISE_KERNEL(MulScalarForward) +REGISTER_HIP_ELEMENTWISE_KERNEL(MulScalarBackward) +REGISTER_HIP_ELEMENTWISE_KERNEL(DivForward) +REGISTER_HIP_ELEMENTWISE_KERNEL(DivBackward) +REGISTER_HIP_ELEMENTWISE_KERNEL(SigmoidForward) +REGISTER_HIP_ELEMENTWISE_KERNEL(SigmoidBackward) + +#undef REGISTER_HIP_ELEMENTWISE_KERNEL diff --git a/infini_train/src/kernels/dcu/embedding.hip b/infini_train/src/kernels/dcu/embedding.hip new file mode 100644 index 00000000..7f2e1e92 --- /dev/null +++ b/infini_train/src/kernels/dcu/embedding.hip @@ -0,0 +1,126 @@ +#include + +#include "infini_train/include/common/dcu/common_dcu.h" +#include "infini_train/include/core/runtime/device_guard.h" +#include "infini_train/include/dispatcher.h" +#include "infini_train/include/tensor.h" + +#include "infini_train/src/core/runtime/dcu/dcu_dispatch.h" +#include "infini_train/src/core/runtime/dcu/dcu_runtime_common.h" + +namespace infini_train::kernels::dcu { + +template +__global__ void EmbeddingForwardKernel(const int64_t *input, T *output, const T *weight, int batch_size, int max_seqlen, + int embed_dim, int vocab_size) { + int idx = (blockIdx.x * blockDim.x + threadIdx.x); + if (idx >= batch_size * max_seqlen * embed_dim) { + return; + } + + int bt = idx / embed_dim; + int b = bt / max_seqlen; + int t = bt % max_seqlen; + int c = idx % embed_dim; + + int ix = static_cast(input[b * max_seqlen + t]); + if (ix < 0 || ix >= vocab_size) { + return; + } + output[b * max_seqlen * embed_dim + t * embed_dim + c] = weight[ix * embed_dim + c]; +} + +std::shared_ptr EmbeddingForward(const std::shared_ptr &input, const std::shared_ptr &weight) { + CHECK(input->Dtype() == DataType::kINT64); + CHECK_EQ(weight->Dims().size(), 2); + + auto device = input->GetDevice(); + const auto &hip_stream = dynamic_cast( + infini_train::core::GetDeviceGuardImpl(device.type())->GetStream(device)) + ->hip_stream(); + + const int batch_size = input->Dims().size() == 2 ? input->Dims()[0] : 1; + const int max_seqlen = input->Dims().size() == 2 ? input->Dims()[1] : input->Dims()[0]; + const int vocab_size = weight->Dims()[0]; + const int embed_dim = weight->Dims()[1]; + auto output_dims = input->Dims(); + output_dims.push_back(embed_dim); + + auto dtype = weight->Dtype(); + auto output = std::make_shared(output_dims, dtype, input->GetDevice()); + int threads_per_block = 256; + int num_blocks = (batch_size * max_seqlen * embed_dim + threads_per_block - 1) / threads_per_block; + + core::dcu::DispatchDcuFunc( + dtype, + [=]() { + EmbeddingForwardKernel<<>>( + static_cast(input->DataPtr()), static_cast(output->DataPtr()), + static_cast(weight->DataPtr()), batch_size, max_seqlen, embed_dim, vocab_size); + }, + "HIP EmbeddingForward"); + + return output; +} + +template +__global__ void EmbeddingBackwardKernel(const int64_t *input_ptr, const T *grad_output_ptr, T *grad_weight_ptr, + int num_tokens, int embedding_dim, int vocab_size) { + int idx = blockIdx.x * blockDim.x + threadIdx.x; + if (idx >= num_tokens) { + return; + } + + int token_id = static_cast(input_ptr[idx]); + if (token_id < 0 || token_id >= vocab_size) { + return; + } + + for (int j = 0; j < embedding_dim; ++j) { + common::dcu::AtomicAdd(&grad_weight_ptr[token_id * embedding_dim + j], + grad_output_ptr[idx * embedding_dim + j]); + } +} + +std::shared_ptr EmbeddingBackward(const std::shared_ptr &input, const std::vector &weight_dims, + const std::shared_ptr &grad_output) { + CHECK(input->Dtype() == DataType::kINT64); + CHECK_EQ(weight_dims.size(), 2); + auto device = input->GetDevice(); + const auto &hip_stream = dynamic_cast( + infini_train::core::GetDeviceGuardImpl(device.type())->GetStream(device)) + ->hip_stream(); + + const int vocab_size = weight_dims[0]; + const int embedding_dim = weight_dims[1]; + CHECK_EQ(input->Dims().size() + 1, grad_output->Dims().size()); + for (int idx = 0; idx < input->Dims().size(); ++idx) { CHECK_EQ(input->Dims()[idx], grad_output->Dims()[idx]); } + CHECK_EQ(*grad_output->Dims().rbegin(), embedding_dim); + + auto dtype = grad_output->Dtype(); + auto grad_weight = std::make_shared(weight_dims, dtype, grad_output->GetDevice()); + const int num_tokens = input->NumElements(); + const int threads_per_block = 256; + const int num_blocks = (num_tokens + threads_per_block - 1) / threads_per_block; + + core::dcu::DispatchDcuFunc( + dtype, + [=]() { + grad_weight->Fill(0.0); + EmbeddingBackwardKernel<<>>( + static_cast(input->DataPtr()), static_cast(grad_output->DataPtr()), + static_cast(grad_weight->DataPtr()), num_tokens, embedding_dim, vocab_size); + }, + "HIP EmbeddingBackward"); + + return grad_weight; +} +} // namespace infini_train::kernels::dcu + +#define REGISTER_HIP_EMBEDDING_KERNEL(kernel_name) \ + REGISTER_KERNEL(infini_train::Device::DeviceType::kDCU, kernel_name, infini_train::kernels::dcu::kernel_name) + +REGISTER_HIP_EMBEDDING_KERNEL(EmbeddingForward) +REGISTER_HIP_EMBEDDING_KERNEL(EmbeddingBackward) + +#undef REGISTER_HIP_EMBEDDING_KERNEL diff --git a/infini_train/src/kernels/dcu/fill.hip b/infini_train/src/kernels/dcu/fill.hip new file mode 100644 index 00000000..dd36afc8 --- /dev/null +++ b/infini_train/src/kernels/dcu/fill.hip @@ -0,0 +1,47 @@ +#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 "infini_train/src/core/runtime/dcu/dcu_dispatch.h" +#include "infini_train/src/core/runtime/dcu/dcu_runtime_common.h" + +namespace infini_train::kernels::dcu { + +template __global__ void FillKernel(T *data, T value, size_t size) { + size_t idx = blockIdx.x * blockDim.x + threadIdx.x; + if (idx < size) { + data[idx] = value; + } +} + +// TODO(dcj): refactor Fill kernel with elementwise template +void Fill(std::shared_ptr tensor, Scalar scalar) { + const int num_tokens = tensor->NumElements(); + const int threads_per_block = 256; + const int num_blocks = (num_tokens + threads_per_block - 1) / threads_per_block; + auto device = tensor->GetDevice(); + const auto &hip_stream = dynamic_cast( + infini_train::core::GetDeviceGuardImpl(device.type())->GetStream(device)) + ->hip_stream(); + + core::dcu::DispatchDcuFunc( + tensor->Dtype(), + [=]() { + const T casted_value = scalar.to(); + FillKernel<<>>(static_cast(tensor->DataPtr()), + casted_value, tensor->NumElements()); + }, + "HIP Fill"); +} +} // namespace infini_train::kernels::dcu + +#define REGISTER_HIP_FILL_KERNEL(kernel_name) \ + REGISTER_KERNEL(infini_train::Device::DeviceType::kDCU, kernel_name, infini_train::kernels::dcu::kernel_name) + +REGISTER_HIP_FILL_KERNEL(Fill) + +#undef REGISTER_HIP_FILL_KERNEL diff --git a/infini_train/src/kernels/dcu/gather.hip b/infini_train/src/kernels/dcu/gather.hip new file mode 100644 index 00000000..d4a4fc2b --- /dev/null +++ b/infini_train/src/kernels/dcu/gather.hip @@ -0,0 +1,232 @@ +#include "glog/logging.h" + +#include "infini_train/include/common/dcu/common_dcu.h" +#include "infini_train/include/core/runtime/device_guard.h" +#include "infini_train/include/dispatcher.h" +#include "infini_train/include/tensor.h" + +#include "infini_train/src/core/runtime/dcu/dcu_dispatch.h" +#include "infini_train/src/core/runtime/dcu/dcu_runtime_common.h" + +namespace infini_train::kernels::dcu { +// FIXME(zbl): This kernel aligns with torch.gather +// Currently named IndexGather to avoid conflict with communication operators +// Should be renamed to Gather later for interface consistency +template +__global__ void IndexGatherForwardKernel(const T *__restrict__ input, const int64_t *__restrict__ norm_index, + T *__restrict__ output, const int64_t *__restrict__ out_dims, + const int64_t *__restrict__ in_strides, + const int64_t *__restrict__ out_strides, int num_dims, int gather_dim, + int64_t dim_size_gather, int64_t total_elements) { + int64_t out_idx = blockIdx.x * blockDim.x + threadIdx.x; + if (out_idx >= total_elements) { + return; + } + + // Normalize like PyTorch: allow negative, clamp to [0, dim_size_gather-1] + int64_t gather_j = norm_index[out_idx]; + gather_j = (gather_j < 0) ? (gather_j + dim_size_gather) : gather_j; + if (gather_j < 0) { + gather_j = 0; + } + if (gather_j >= dim_size_gather) { + gather_j = dim_size_gather - 1; + } + + int64_t in_linear = 0, tmp = out_idx; +#pragma unroll + for (int d = 0; d < num_dims; ++d) { + int64_t coord = tmp / out_strides[d]; + tmp -= coord * out_strides[d]; + in_linear += ((d == gather_dim) ? gather_j : coord) * in_strides[d]; + } + output[out_idx] = input[in_linear]; +} + +std::shared_ptr IndexGatherForward(const std::shared_ptr &input, const std::shared_ptr &index, + int64_t dim) { + const auto &in_dims = input->Dims(); + const auto &idx_dims = index->Dims(); + CHECK_EQ(in_dims.size(), idx_dims.size()); + CHECK(input->GetDevice().type() == index->GetDevice().type()); + CHECK(input->GetDevice().index() == index->GetDevice().index()); + + const int64_t num_dims = in_dims.size(); + if (dim < 0) { + dim += num_dims; + } + CHECK_GE(dim, 0); + CHECK_LT(dim, num_dims); + + // NOTE(zbl): Assume index to be int64 Tensors + CHECK(index->Dtype() == DataType::kINT64); + + for (int d = 0; d < num_dims; ++d) { + if (d == dim) { + continue; + } + // Align with PyTorch semantics: index.size(d) <= input.size(d) for d != dim + CHECK_LE(idx_dims[d], in_dims[d]) + << "index.size(" << d << ") must be <= input.size(" << d << ") on non-gather dims"; + } + + const auto device = input->GetDevice(); + const auto &stream = dynamic_cast( + infini_train::core::GetDeviceGuardImpl(device.type())->GetStream(device)) + ->hip_stream(); + + auto dtype = input->Dtype(); + auto out = std::make_shared(idx_dims, dtype, device); + + auto in_strides = ComputeStrides(in_dims); + auto out_strides = ComputeStrides(idx_dims); + const int64_t total_elements = index->NumElements(); + + const int64_t gather_dim_size = in_dims[dim]; + + int64_t *dev_buf = nullptr; + HIP_CHECK(hipMallocAsync(&dev_buf, (3 * num_dims) * sizeof(int64_t), stream)); + int64_t *out_dims_dev = dev_buf + 0 * num_dims; + int64_t *in_strides_dev = dev_buf + 1 * num_dims; + int64_t *out_strides_dev = dev_buf + 2 * num_dims; + + HIP_CHECK( + hipMemcpyAsync(out_dims_dev, idx_dims.data(), num_dims * sizeof(int64_t), hipMemcpyHostToDevice, stream)); + HIP_CHECK( + hipMemcpyAsync(in_strides_dev, in_strides.data(), num_dims * sizeof(int64_t), hipMemcpyHostToDevice, stream)); + HIP_CHECK(hipMemcpyAsync(out_strides_dev, out_strides.data(), num_dims * sizeof(int64_t), hipMemcpyHostToDevice, + stream)); + + const int threads = 256; + const int blocks = (total_elements + threads - 1) / threads; + + core::dcu::DispatchDcuFunc( + dtype, + [=]() { + IndexGatherForwardKernel<<>>( + static_cast(input->DataPtr()), static_cast(index->DataPtr()), + static_cast(out->DataPtr()), out_dims_dev, in_strides_dev, out_strides_dev, (int)num_dims, + (int)dim, gather_dim_size, total_elements); + }, + "HIP IndexGatherForward"); + + HIP_CHECK(hipFree(dev_buf)); + return out; +} + +template +__global__ void IndexGatherBackwardKernel(const T *__restrict__ grad_output, const int64_t *__restrict__ index, + T *__restrict__ grad_input, const int64_t *__restrict__ out_dims, + const int64_t *__restrict__ in_strides, + const int64_t *__restrict__ out_strides, int num_dims, int gather_dim, + int64_t dim_size_gather, int64_t total_elements) { + int64_t out_idx = blockIdx.x * blockDim.x + threadIdx.x; + if (out_idx >= total_elements) { + return; + } + + int64_t gather_j = index[out_idx]; + gather_j = (gather_j < 0) ? (gather_j + dim_size_gather) : gather_j; + if (gather_j < 0) { + gather_j = 0; + } + if (gather_j >= dim_size_gather) { + gather_j = dim_size_gather - 1; + } + + int64_t in_linear = 0; + int64_t tmp = out_idx; +#pragma unroll + for (int d = 0; d < num_dims; ++d) { + int64_t coord = tmp / out_strides[d]; + tmp -= coord * out_strides[d]; + if (d == gather_dim) { + in_linear += gather_j * in_strides[d]; + } else { + in_linear += coord * in_strides[d]; + } + } + common::dcu::AtomicAdd(&grad_input[in_linear], grad_output[out_idx]); +} + +std::shared_ptr IndexGatherBackward(const std::shared_ptr &grad_output, + const std::shared_ptr &index, int64_t dim, + const std::vector &input_dims) { + const auto &in_dims = input_dims; + const auto &idx_dims = index->Dims(); + CHECK_EQ(in_dims.size(), idx_dims.size()); + const int64_t num_dims = in_dims.size(); + if (dim < 0) { + dim += num_dims; + } + CHECK_GE(dim, 0); + CHECK_LT(dim, num_dims); + + // NOTE(zbl): Assume index to be int64 Tensors + CHECK(index->Dtype() == DataType::kINT64); + + for (int d = 0; d < num_dims; ++d) { + if (d == dim) { + continue; + } + CHECK_EQ(in_dims[d], idx_dims[d]); + } + + auto dtype = grad_output->Dtype(); + auto grad_input = std::make_shared(in_dims, dtype, grad_output->GetDevice()); + grad_input->Fill(0.0); + + auto in_strides = ComputeStrides(in_dims); + auto out_strides = ComputeStrides(idx_dims); + const int64_t total_elements + = std::accumulate(idx_dims.begin(), idx_dims.end(), (int64_t)1, std::multiplies{}); + const int64_t gather_dim_size = in_dims[dim]; + + int64_t *dev_buf = nullptr; + const size_t n_out = idx_dims.size(); + const size_t n_in_strides = in_dims.size(); + const size_t n_out_strides = idx_dims.size(); + const size_t total_i64 = n_out + n_in_strides + n_out_strides; + + auto device = grad_output->GetDevice(); + const auto &stream = dynamic_cast( + infini_train::core::GetDeviceGuardImpl(device.type())->GetStream(device)) + ->hip_stream(); + + HIP_CHECK(hipMallocAsync(&dev_buf, total_i64 * sizeof(int64_t), stream)); + int64_t *out_dims_dev = dev_buf; + int64_t *in_strides_dev = out_dims_dev + n_out; + int64_t *out_strides_dev = in_strides_dev + n_in_strides; + + HIP_CHECK(hipMemcpyAsync(out_dims_dev, idx_dims.data(), n_out * sizeof(int64_t), hipMemcpyHostToDevice, stream)); + HIP_CHECK(hipMemcpyAsync(in_strides_dev, in_strides.data(), n_in_strides * sizeof(int64_t), + hipMemcpyHostToDevice, stream)); + HIP_CHECK(hipMemcpyAsync(out_strides_dev, out_strides.data(), n_out_strides * sizeof(int64_t), + hipMemcpyHostToDevice, stream)); + + const int threads = 256; + const int blocks = (int)((total_elements + threads - 1) / threads); + + core::dcu::DispatchDcuFunc( + dtype, + [=]() { + IndexGatherBackwardKernel<<>>( + static_cast(grad_output->DataPtr()), static_cast(index->DataPtr()), + static_cast(grad_input->DataPtr()), out_dims_dev, in_strides_dev, out_strides_dev, (int)num_dims, + (int)dim, gather_dim_size, total_elements); + }, + "HIP IndexGatherBackward"); + + HIP_CHECK(hipFree(dev_buf)); + return grad_input; +} + +} // namespace infini_train::kernels::dcu + +#define REGISTER_HIP_GATHER_KERNEL(kernel_name) \ + REGISTER_KERNEL(infini_train::Device::DeviceType::kDCU, kernel_name, infini_train::kernels::dcu::kernel_name) + +REGISTER_HIP_GATHER_KERNEL(IndexGatherForward) +REGISTER_HIP_GATHER_KERNEL(IndexGatherBackward) + +#undef REGISTER_HIP_GATHER_KERNEL diff --git a/infini_train/src/kernels/dcu/layernorm.hip b/infini_train/src/kernels/dcu/layernorm.hip new file mode 100644 index 00000000..3f0287e8 --- /dev/null +++ b/infini_train/src/kernels/dcu/layernorm.hip @@ -0,0 +1,208 @@ +#include + +#include "infini_train/include/common/dcu/common_dcu.h" +#include "infini_train/include/common/dcu/kernel_helper.cuh" +#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 "infini_train/src/core/runtime/dcu/dcu_dispatch.h" +#include "infini_train/src/core/runtime/dcu/dcu_runtime_common.h" + +namespace infini_train::kernels::dcu { + +template +__global__ void LayerNormForwardKernel(const T *input, const T *weight, const T *bias, float *mean_out, float *rstd_out, + T *output, float eps, int embed_dim) { + using BlockReduce = hipcub::BlockReduce; + __shared__ typename BlockReduce::TempStorage temp_storage_mean; + __shared__ typename BlockReduce::TempStorage temp_storage_rstd; + __shared__ float shared_mean; + __shared__ float shared_rstd; + + const int token_idx = blockIdx.x; + const T *x = input + token_idx * embed_dim; + T *y = output + token_idx * embed_dim; + + float sum = 0.0f; + float sqsum = 0.0f; + + for (int i = threadIdx.x; i < embed_dim; i += BLOCK_SIZE) { + float val = common::dcu::Cast(x[i]); + sum += val; + sqsum += val * val; + } + + float total_sum = BlockReduce(temp_storage_mean).Sum(sum); + float total_sqsum = BlockReduce(temp_storage_rstd).Sum(sqsum); + + if (threadIdx.x == 0) { + float mean = total_sum / embed_dim; + float var = total_sqsum / embed_dim - mean * mean; + float rstd = rsqrtf(var + eps); + shared_mean = mean; + shared_rstd = rstd; + if (mean_out) { + mean_out[token_idx] = mean; + } + if (rstd_out) { + rstd_out[token_idx] = rstd; + } + } + __syncthreads(); + + for (int i = threadIdx.x; i < embed_dim; i += BLOCK_SIZE) { + float norm = (common::dcu::Cast(x[i]) - shared_mean) * shared_rstd; + y[i] = common::dcu::Cast(norm * common::dcu::Cast(weight[i]) + common::dcu::Cast(bias[i])); + } +} + +std::tuple, std::shared_ptr, std::shared_ptr> +LayerNormForward(const std::shared_ptr &input, const std::shared_ptr &weight, + const std::shared_ptr &bias, const float eps) { + CHECK_EQ(input->Dims().size(), 3); + CHECK_LE(input->Dims()[2], weight->Dims()[0]); + CHECK_LE(input->Dims()[2], bias->Dims()[0]); + + const int batch_size = input->Dims()[0]; + const int max_seqlen = input->Dims()[1]; + const int embed_dim = input->Dims()[2]; + + auto dtype = input->Dtype(); + + auto output = std::make_shared(input->Dims(), dtype, input->GetDevice()); + auto mean = std::make_shared(std::vector{batch_size, max_seqlen}, DataType::kFLOAT32, + input->GetDevice()); + auto rstd = std::make_shared(std::vector{batch_size, max_seqlen}, DataType::kFLOAT32, + input->GetDevice()); + + constexpr int BLOCK_SIZE = 256; + int threads_per_block = BLOCK_SIZE; + int num_blocks = batch_size * max_seqlen; + + auto device = input->GetDevice(); + const auto &hip_stream = dynamic_cast( + infini_train::core::GetDeviceGuardImpl(device.type())->GetStream(device)) + ->hip_stream(); + + core::dcu::DispatchDcuFunc( + dtype, + [=]() { + mean->Fill(0.0); + rstd->Fill(0.0); + LayerNormForwardKernel<<>>( + static_cast(input->DataPtr()), static_cast(weight->DataPtr()), + static_cast(bias->DataPtr()), static_cast(mean->DataPtr()), + static_cast(rstd->DataPtr()), static_cast(output->DataPtr()), eps, embed_dim); + }, + "HIP LayerNormForward"); + + return {output, mean, rstd}; +} + +template +__global__ void LayerNormBackwardKernel(const T *__restrict__ input, const T *__restrict__ grad_output, + const float *__restrict__ mean, const float *__restrict__ rstd, + const T *__restrict__ weight, T *__restrict__ grad_input, + T *__restrict__ grad_weight, T *__restrict__ grad_bias, int embed_dim, + size_t weight_num_elements, size_t bias_num_elements) { + using BlockReduce = hipcub::BlockReduce; + __shared__ typename BlockReduce::TempStorage temp_storage_mean; + __shared__ typename BlockReduce::TempStorage temp_storage_norm; + __shared__ float shared_mean; + __shared__ float shared_norm; + + int tid = threadIdx.x; + int token_idx = blockIdx.x; + + const T *input_ptr = input + token_idx * embed_dim; + const T *grad_output_ptr = grad_output + token_idx * embed_dim; + T *grad_input_ptr = grad_input + token_idx * embed_dim; + + float mean_val = mean[token_idx]; + float rstd_val = rstd[token_idx]; + + float dnorm_mean = 0.f; + float dnorm_norm_mean = 0.f; + + for (int i = tid; i < embed_dim; i += BLOCK_SIZE) { + float dnorm = common::dcu::Cast(common::dcu::Mul(weight[i], grad_output_ptr[i])); + dnorm_mean += dnorm; + dnorm_norm_mean += dnorm * (common::dcu::Cast(input_ptr[i]) - mean_val); + } + + dnorm_mean = BlockReduce(temp_storage_mean).Sum(dnorm_mean); + dnorm_norm_mean = BlockReduce(temp_storage_norm).Sum(dnorm_norm_mean); + + if (tid == 0) { + float mean_d = dnorm_mean / embed_dim; + float norm_d = (dnorm_norm_mean / embed_dim) * rstd_val - mean_d * mean_val * rstd_val; + shared_mean = mean_d; + shared_norm = norm_d; + } + __syncthreads(); + + for (int i = tid; i < embed_dim; i += BLOCK_SIZE) { + float norm = (common::dcu::Cast(input_ptr[i]) - mean_val) * rstd_val; + float grad_output_val = common::dcu::Cast(grad_output_ptr[i]); + + grad_input_ptr[i] = common::dcu::Cast( + (common::dcu::Cast(weight[i]) * grad_output_val - shared_mean - norm * shared_norm) * rstd_val); + + common::dcu::fastAtomicAdd(grad_weight, i, weight_num_elements, + common::dcu::Cast(grad_output_val * norm), true); + common::dcu::fastAtomicAdd(grad_bias, i, bias_num_elements, grad_output_ptr[i], true); + } +} + +std::tuple, std::shared_ptr, std::shared_ptr> +LayerNormBackward(const std::shared_ptr &input, const std::shared_ptr &weight, + const std::shared_ptr &bias, const std::shared_ptr &mean, + const std::shared_ptr &rstd, const std::shared_ptr &grad_output) { + const int batch_size = input->Dims()[0]; + const int max_seqlen = input->Dims()[1]; + const int embed_dim = input->Dims()[2]; + + auto dtype = input->Dtype(); + CHECK(dtype == weight->Dtype() && dtype == bias->Dtype() && dtype == grad_output->Dtype() + && mean->Dtype() == DataType::kFLOAT32 && rstd->Dtype() == DataType::kFLOAT32); + + auto grad_input = std::make_shared(input->Dims(), dtype, grad_output->GetDevice()); + auto grad_weight = std::make_shared(weight->Dims(), dtype, grad_output->GetDevice()); + auto grad_bias = std::make_shared(bias->Dims(), dtype, grad_output->GetDevice()); + + constexpr int BLOCK_SIZE = 256; + int threads_per_block = BLOCK_SIZE; + int num_blocks = batch_size * max_seqlen; + + auto device = input->GetDevice(); + const auto &hip_stream = dynamic_cast( + infini_train::core::GetDeviceGuardImpl(device.type())->GetStream(device)) + ->hip_stream(); + core::dcu::DispatchDcuFunc( + dtype, + [=]() { + grad_input->Fill(0.0); + grad_weight->Fill(0.0); + grad_bias->Fill(0.0); + LayerNormBackwardKernel<<>>( + static_cast(input->DataPtr()), static_cast(grad_output->DataPtr()), + static_cast(mean->DataPtr()), static_cast(rstd->DataPtr()), + static_cast(weight->DataPtr()), static_cast(grad_input->DataPtr()), + static_cast(grad_weight->DataPtr()), static_cast(grad_bias->DataPtr()), embed_dim, + grad_weight->NumElements(), grad_bias->NumElements()); + }, + "HIP LayerNormBackward"); + + return {grad_input, grad_weight, grad_bias}; +} +} // namespace infini_train::kernels::dcu + +#define REGISTER_HIP_LAYERNORM_KERNEL(kernel_name) \ + REGISTER_KERNEL(infini_train::Device::DeviceType::kDCU, kernel_name, infini_train::kernels::dcu::kernel_name) + +REGISTER_HIP_LAYERNORM_KERNEL(LayerNormForward) +REGISTER_HIP_LAYERNORM_KERNEL(LayerNormBackward) + +#undef REGISTER_HIP_LAYERNORM_KERNEL diff --git a/infini_train/src/kernels/dcu/linear.hip b/infini_train/src/kernels/dcu/linear.hip new file mode 100644 index 00000000..f72d41cb --- /dev/null +++ b/infini_train/src/kernels/dcu/linear.hip @@ -0,0 +1,504 @@ +#include +#include +#include +#include + +#include +#include + +#include "infini_train/include/autograd/linear.h" +#include "infini_train/include/common/dcu/common_dcu.h" +#include "infini_train/include/common/dcu/kernel_helper.cuh" +#include "infini_train/include/core/runtime/device_guard.h" +#include "infini_train/include/dispatcher.h" +#include "infini_train/include/tensor.h" + +#include "infini_train/src/core/runtime/dcu/dcu_dispatch.h" +#include "infini_train/src/core/runtime/dcu/dcu_runtime_common.h" + +namespace infini_train::kernels::dcu { + +std::shared_ptr MatmulForward(const std::shared_ptr &input, const std::shared_ptr &other) { + /* + output[*, m, n] = input[*, m, k] * other[*, k, n] + */ + const auto &input_dims = input->Dims(); + const auto &other_dims = other->Dims(); + + CHECK_GE(input_dims.size(), 2); + CHECK_GE(other_dims.size(), 2); + CHECK_EQ(input_dims.size(), other_dims.size()); + + const int64_t m = input_dims[input_dims.size() - 2]; + const int64_t k = input_dims[input_dims.size() - 1]; + CHECK_EQ(k, other_dims[other_dims.size() - 2]); + const int64_t n = other_dims[other_dims.size() - 1]; + + const int64_t bs = std::accumulate(input_dims.rbegin() + 2, input_dims.rend(), 1, std::multiplies{}); + for (int64_t i = 0; i < input_dims.size() - 2; ++i) { + CHECK_EQ(input_dims[i], other_dims[i]) << "Batch dims must match"; + } + + auto dtype = input->Dtype(); + std::vector output_dims = input_dims; + output_dims[output_dims.size() - 1] = n; + auto output = std::make_shared(output_dims, dtype, input->GetDevice()); + + auto device = input->GetDevice(); + const float alpha = 1.0f, beta = 0.0f; + hipblasHandle_t handle = dynamic_cast( + infini_train::core::GetDeviceGuardImpl(device.type())->GetBlasHandle(device)) + ->hipblas_handle(); + + // cuBLAS is colmun-major + // output = input * other --> output.T = other.T * input.T + // C = A * B ==> output.T[*, n, m] = other.T[*, n, k] * input.T[*, k, m] + // C = output.T[*, n, m] + // A = other.T[*, n, k] + // B = input.T[*, k, m] + int lda = n; + int ldb = k; + int ldc = n; + int64_t stride_a = n * k; + int64_t stride_b = k * m; + int64_t stride_c = m * n; + // NOTE(zbl): the last hipblasGemmAlgo_t param has no effect on GPU arch >= sm_80(Ampere) + + switch (dtype) { + DISPATCH_CASE(WRAP(HIPBLAS_CHECK(hipblasGemmStridedBatchedEx( + handle, HIPBLAS_OP_N, HIPBLAS_OP_N, n, m, k, &alpha, other->DataPtr(), HIPBLAS_R_32F, lda, + stride_a, input->DataPtr(), HIPBLAS_R_32F, ldb, stride_b, &beta, output->DataPtr(), HIPBLAS_R_32F, + ldc, stride_c, bs, HIPBLAS_R_32F, HIPBLAS_GEMM_DEFAULT));), + DataType::kFLOAT32) + DISPATCH_CASE(WRAP(HIPBLAS_CHECK(hipblasGemmStridedBatchedEx( + handle, HIPBLAS_OP_N, HIPBLAS_OP_N, n, m, k, &alpha, other->DataPtr(), HIPBLAS_R_16B, lda, + stride_a, input->DataPtr(), HIPBLAS_R_16B, ldb, stride_b, &beta, output->DataPtr(), HIPBLAS_R_16B, + ldc, stride_c, bs, HIPBLAS_R_32F, HIPBLAS_GEMM_DEFAULT));), + DataType::kBFLOAT16) + default: + LOG_UNSUPPORTED_DTYPE(dtype, "HIP MatmulForward"); + } + + return output; +} + +std::tuple, std::shared_ptr> +MatmulBackward(const std::shared_ptr &input, const std::shared_ptr &other, + const std::shared_ptr &grad_output) { + /* + grad_input[*, m, k] = grad_output[*, m, n] * other[*, k, n]^T + grad_other[*, k, n] = input[*, m, k]^T * grad_output[*, m, n] + */ + + auto input_dtype = input->Dtype(); + auto other_dtype = other->Dtype(); + auto grad_output_dtype = grad_output->Dtype(); + // Compute dtype determined by saved tensors (forward compute dtype), not grad_output + DataType compute_dtype = PromoteDataTypes(input_dtype, other_dtype); + + auto input_promoted = input_dtype == compute_dtype ? input : std::make_shared(input->To(compute_dtype)); + auto other_promoted = other_dtype == compute_dtype ? other : std::make_shared(other->To(compute_dtype)); + auto grad_output_promoted + = grad_output_dtype == compute_dtype ? grad_output : std::make_shared(grad_output->To(compute_dtype)); + + const auto &input_dims = input->Dims(); + const auto &other_dims = other->Dims(); + const auto &grad_output_dims = grad_output->Dims(); + + CHECK_GE(input_dims.size(), 2); + CHECK_EQ(input_dims.size(), other_dims.size()); + CHECK_EQ(input_dims.size(), grad_output_dims.size()); + + const int64_t m = input_dims[input_dims.size() - 2]; + const int64_t k = input_dims[input_dims.size() - 1]; + const int64_t n = other_dims[other_dims.size() - 1]; + CHECK_EQ(k, other_dims[other_dims.size() - 2]); + CHECK_EQ(m, grad_output_dims[grad_output_dims.size() - 2]); + CHECK_EQ(n, grad_output_dims[grad_output_dims.size() - 1]); + + const int64_t bs = std::accumulate(input_dims.rbegin() + 2, input_dims.rend(), 1, std::multiplies{}); + for (int64_t i = 0; i < input_dims.size() - 2; ++i) { + CHECK_EQ(input_dims[i], other_dims[i]) << "Batch dims must match"; + CHECK_EQ(input_dims[i], grad_output_dims[i]) << "Batch dims must match"; + } + + // For bf16 compute, output in fp32 to preserve accumulation precision (matches PyTorch behavior) + auto output_dtype = (compute_dtype == DataType::kBFLOAT16) ? DataType::kFLOAT32 : compute_dtype; + auto grad_input = std::make_shared(input_dims, output_dtype, grad_output->GetDevice()); + auto grad_other = std::make_shared(other_dims, output_dtype, grad_output->GetDevice()); + + // No Fill(0) needed: cuBLAS beta=0.0f means C is fully overwritten, never read. + + auto device = input_promoted->GetDevice(); + const float alpha = 1.0f, beta = 0.0f; + hipblasHandle_t handle = dynamic_cast( + infini_train::core::GetDeviceGuardImpl(device.type())->GetBlasHandle(device)) + ->hipblas_handle(); + + { + // cuBLAS is colmun-major + // grad_input = grad_output * other.T --> grad_input.T = other * grad_output.T + // C = A.T * B ==> grad_input.T[*, k, m] = other[*, k, n] * grad_output.T[*, n, m] + // C = grad_input.T[*, k, m] + // A = other.T[*, n, k] + // B = grad_output.T[*, n, m] + const int lda = n, ldb = n, ldc = k; + const int64_t stride_a = k * n; + const int64_t stride_b = n * m; + const int64_t stride_c = m * k; + switch (compute_dtype) { + DISPATCH_CASE(WRAP(HIPBLAS_CHECK(hipblasGemmStridedBatchedEx( + handle, HIPBLAS_OP_T, HIPBLAS_OP_N, k, m, n, &alpha, other_promoted->DataPtr(), HIPBLAS_R_32F, + lda, stride_a, grad_output_promoted->DataPtr(), HIPBLAS_R_32F, ldb, stride_b, &beta, + grad_input->DataPtr(), HIPBLAS_R_32F, ldc, stride_c, bs, HIPBLAS_R_32F, HIPBLAS_GEMM_DEFAULT));), + DataType::kFLOAT32) + DISPATCH_CASE(WRAP(HIPBLAS_CHECK(hipblasGemmStridedBatchedEx( + handle, HIPBLAS_OP_T, HIPBLAS_OP_N, k, m, n, &alpha, other_promoted->DataPtr(), HIPBLAS_R_16B, + lda, stride_a, grad_output_promoted->DataPtr(), HIPBLAS_R_16B, ldb, stride_b, &beta, + grad_input->DataPtr(), HIPBLAS_R_32F, ldc, stride_c, bs, HIPBLAS_R_32F, HIPBLAS_GEMM_DEFAULT));), + DataType::kBFLOAT16) + } + } + + { + // cuBLAS is colmun-major + // grad_other = input.T * grad_output --> grad_other.T = grad_output.T * input + // C = A * B.T ==> grad_other.T[*, n, k] = grad_output.T[*, n, m] * input[*, m, k] + // C = grad_other.T[*, n, k] + // A = grad_output.T[*, n, m] + // B = input.T[*, k, m] + const int lda = n, ldb = k, ldc = n; + const int64_t stride_a = n * m; + const int64_t stride_b = k * m; + const int64_t stride_c = n * k; + switch (compute_dtype) { + DISPATCH_CASE(WRAP(HIPBLAS_CHECK(hipblasGemmStridedBatchedEx( + handle, HIPBLAS_OP_N, HIPBLAS_OP_T, n, k, m, &alpha, grad_output_promoted->DataPtr(), + HIPBLAS_R_32F, lda, stride_a, input_promoted->DataPtr(), HIPBLAS_R_32F, ldb, stride_b, &beta, + grad_other->DataPtr(), HIPBLAS_R_32F, ldc, stride_c, bs, HIPBLAS_R_32F, HIPBLAS_GEMM_DEFAULT));), + DataType::kFLOAT32) + DISPATCH_CASE(WRAP(HIPBLAS_CHECK(hipblasGemmStridedBatchedEx( + handle, HIPBLAS_OP_N, HIPBLAS_OP_T, n, k, m, &alpha, grad_output_promoted->DataPtr(), + HIPBLAS_R_16B, lda, stride_a, input_promoted->DataPtr(), HIPBLAS_R_16B, ldb, stride_b, &beta, + grad_other->DataPtr(), HIPBLAS_R_32F, ldc, stride_c, bs, HIPBLAS_R_32F, HIPBLAS_GEMM_DEFAULT));), + DataType::kBFLOAT16) + } + } + + return {grad_input, grad_other}; +} + +template __global__ void BiasCopyKernel(T *output, const T *bias, int bs, int out_features) { + int idx = blockIdx.x * blockDim.x + threadIdx.x; + if (idx >= bs * out_features) { + return; + } + int j = idx % out_features; + output[idx] = bias[j]; +} + +std::shared_ptr LinearForward(const std::shared_ptr &input, const std::shared_ptr &weight, + bool transpose, const std::shared_ptr &bias) { + + /* + !transpose: output = input * weight + bias + output[*, out_features] = input[*, in_features] * weight[in_features, out_features] + bias[out_features] + + transpose: output = input * weight^T + bias + output[*, out_features] = input[*, in_features] * weight[out_features, in_features]^T + bias[out_features] + */ + + const auto &input_dims = input->Dims(); + CHECK_GE(input_dims.size(), 2); + const int64_t bs = std::accumulate(input_dims.rbegin() + 1, input_dims.rend(), 1, std::multiplies{}); + const int64_t in_features = *input_dims.rbegin(); + + const auto &weight_dims = weight->Dims(); + CHECK_EQ(weight_dims.size(), 2); + CHECK_EQ(in_features, weight_dims[transpose ? 1 : 0]); + + // As for hipblas: + // C = alpha * op(B) * op(A) + beta * C + // Dimensions: + // input: (bs, in_features) + // weight: (in_features, out_features) or (out_features, in_features) if transposed + // output: (bs, out_features) + const int64_t out_features = weight_dims[transpose ? 0 : 1]; + + auto dtype = input->Dtype(); + auto output_dims = input_dims; + *output_dims.rbegin() = out_features; + auto output = std::make_shared(output_dims, dtype, input->GetDevice()); + + auto device = input->GetDevice(); + const auto &hip_stream = dynamic_cast( + infini_train::core::GetDeviceGuardImpl(device.type())->GetStream(device)) + ->hip_stream(); + + if (bias) { + CHECK_EQ(bias->Dims().size(), 1); + CHECK_EQ(bias->Dims()[0], out_features); + int threads_per_block = 256; + int num_blocks = (bs * out_features + threads_per_block - 1) / threads_per_block; + + core::dcu::DispatchDcuFunc( + dtype, + [=]() { + BiasCopyKernel<<>>( + static_cast(output->DataPtr()), static_cast(bias->DataPtr()), bs, out_features); + }, + "HIP LinearForward"); + } else { + output->Fill(0.0); + } + + const float alpha = 1.0f; + const float beta = 1.0f; + auto trans_a = transpose ? HIPBLAS_OP_T : HIPBLAS_OP_N; + auto trans_b = HIPBLAS_OP_N; + auto lda = transpose ? in_features : out_features; + hipblasHandle_t handle = dynamic_cast( + infini_train::core::GetDeviceGuardImpl(device.type())->GetBlasHandle(device)) + ->hipblas_handle(); + + // TODO(zbl): use hipblasSgemv if possible for convenience and simplicity + // + // - if a is transposed: + // weight is [out_features, in_features] here + // output = input * weight.T --> output.T = weight * input.T + // C = output.T[out_features, bs] + // A = weight.T[in_features, out_features] + // B = input.T[in_features, bs] + // + // - if a is not transposed: + // output = input * weight --> output.T = weight.T * input.T + // C = output.T[out_features, bs] + // A = weight.T[out_features, in_features] + // B = input.T[in_features, bs] + switch (input->Dtype()) { + DISPATCH_CASE(WRAP({ + HIPBLAS_CHECK(hipblasSgemm(handle, trans_a, trans_b, out_features, bs, in_features, &alpha, + static_cast(weight->DataPtr()), lda, + static_cast(input->DataPtr()), in_features, &beta, + static_cast(output->DataPtr()), out_features)); + }), + DataType::kFLOAT32) + DISPATCH_CASE(WRAP({ + HIPBLAS_CHECK(hipblasGemmEx(handle, trans_a, trans_b, out_features, bs, in_features, &alpha, + weight->DataPtr(), HIPBLAS_R_16B, lda, input->DataPtr(), HIPBLAS_R_16B, + in_features, &beta, output->DataPtr(), HIPBLAS_R_16B, out_features, + HIPBLAS_R_32F, HIPBLAS_GEMM_DEFAULT)); + }), + DataType::kBFLOAT16) + } + + return output; +} + +template +__global__ void ReduceColumnsKernel(const TIn *__restrict__ input, TOut *__restrict__ output, int num_rows, + int num_cols) { + using BlockReduce = hipcub::BlockReduce; + __shared__ typename BlockReduce::TempStorage temp_storage; + + int row = blockIdx.x; + float sum = 0.0f; + + for (int col = threadIdx.x; col < num_cols; col += blockDim.x) { + sum += common::dcu::Cast(input[row * num_cols + col]); + } + + float reduced = BlockReduce(temp_storage).Sum(sum); + + if (threadIdx.x == 0) { + output[row] = reduced; + } +} + +std::tuple, std::shared_ptr, std::shared_ptr> +LinearBackward(const std::shared_ptr &input, const std::shared_ptr &weight, bool transpose, + int64_t in_features, int64_t out_features, const std::vector &input_dims, + const std::shared_ptr &grad_output, bool bias, + infini_train::autograd::LinearGradFlags grad_flags) { + const auto compute_grad_input = grad_flags.input; + const auto compute_grad_weight = grad_flags.weight; + const auto compute_grad_bias = grad_flags.bias; + + CHECK_GE(input_dims.size(), 2); + const int64_t bs = std::accumulate(input_dims.rbegin() + 1, input_dims.rend(), 1, std::multiplies{}); + + const std::vector weight_dims + = transpose ? std::vector{out_features, in_features} : std::vector{in_features, out_features}; + + auto dtype = grad_output->Dtype(); + + // For type promotion, use available tensors + DataType input_dtype = input ? input->Dtype() : (weight ? weight->Dtype() : dtype); + DataType weight_dtype = weight ? weight->Dtype() : (input ? input->Dtype() : dtype); + // Compute dtype determined by saved tensors (forward compute dtype), not grad_output + DataType compute_dtype = PromoteDataTypes(input_dtype, weight_dtype); + + auto grad_output_promoted + = dtype == compute_dtype ? grad_output : std::make_shared(grad_output->To(compute_dtype)); + + // For bf16 compute, accumulate in fp32 to preserve precision (matches PyTorch behavior). + auto output_dtype = (compute_dtype == DataType::kBFLOAT16) ? DataType::kFLOAT32 : compute_dtype; + + // Allocate only needed gradient tensors (selective save: input/weight may be nullptr). + std::shared_ptr grad_input = nullptr; + std::shared_ptr grad_weight = nullptr; + std::shared_ptr grad_bias = nullptr; + + if (compute_grad_input) { + grad_input = std::make_shared(input_dims, output_dtype, grad_output->GetDevice()); + } + if (compute_grad_weight) { + grad_weight = std::make_shared(weight_dims, output_dtype, grad_output->GetDevice()); + } + // No Fill(0) needed: cuBLAS beta=0.0f fully overwrites output, and ReduceColumnsKernel assigns directly. + if (compute_grad_bias && bias) { + grad_bias + = std::make_shared(std::vector{out_features}, output_dtype, grad_output->GetDevice()); + } + + auto device = grad_output->GetDevice(); + const auto &hip_stream = dynamic_cast( + infini_train::core::GetDeviceGuardImpl(device.type())->GetStream(device)) + ->hip_stream(); + + float alpha = 1.0f; + float beta = 0.0f; + + hipblasHandle_t handle = dynamic_cast( + infini_train::core::GetDeviceGuardImpl(device.type())->GetBlasHandle(device)) + ->hipblas_handle(); + + switch (compute_dtype) { + // TODO(zbl): use hipblasSgemv if possible + DISPATCH_CASE( + WRAP({ + if (compute_grad_input) { + // - if transpose: + // weight is [out_features, in_features] here + // d_input = d_output * weight --> d_input.T = weight.T * d_output.T + // C = d_input.T[in_features, bs] + // A = weight.T[in_features, out_features] + // B = d_output.T[out_features, bs] + // + // - if not transpose: + // weight is [in_features, out_features] here + // d_input = d_output * weight.T --> d_input.T = weight * d_output.T + // C = d_input.T[in_features, bs] + // A = weight.T[out_features, in_features] + // B = d_output.T[out_features, bs] + CHECK(weight != nullptr) + << "compute_grad_input=true but weight is nullptr (selective save mismatch)"; + auto weight_promoted + = weight_dtype == compute_dtype ? weight : std::make_shared(weight->To(compute_dtype)); + auto trans_a1 = transpose ? HIPBLAS_OP_N : HIPBLAS_OP_T; + auto lda1 = transpose ? in_features : out_features; + HIPBLAS_CHECK(hipblasSgemm(handle, trans_a1, HIPBLAS_OP_N, in_features, bs, out_features, &alpha, + static_cast(weight_promoted->DataPtr()), lda1, + static_cast(grad_output_promoted->DataPtr()), out_features, + &beta, static_cast(grad_input->DataPtr()), in_features)); + } + if (compute_grad_weight) { + // - if transpose: + // d_weight = d_output.T * input --> d_weight.T = input.T * d_output + // C = d_weight.T[in_features, out_features] + // A = input.T[in_features, bs] + // B = d_output.T[out_features, bs] + // + // - if not transpose: + // d_weight = input.T * d_output --> d_weight.T = d_output.T * input + // C = d_weight.T[out_features, in_features] + // A = d_output.T[out_features, bs] + // B = input.T[in_features, bs] + CHECK(input != nullptr) + << "compute_grad_weight=true but input is nullptr (selective save mismatch)"; + auto input_promoted + = input_dtype == compute_dtype ? input : std::make_shared(input->To(compute_dtype)); + auto trans_a2 = HIPBLAS_OP_N; + auto trans_b2 = HIPBLAS_OP_T; + int m2 = transpose ? in_features : out_features; + int n2 = transpose ? out_features : in_features; + const void *a2 = transpose ? input_promoted->DataPtr() : grad_output_promoted->DataPtr(); + const void *b2 = transpose ? grad_output_promoted->DataPtr() : input_promoted->DataPtr(); + auto lda2 = transpose ? in_features : out_features; + auto ldb2 = transpose ? out_features : in_features; + auto ldc2 = transpose ? in_features : out_features; + HIPBLAS_CHECK(hipblasSgemm(handle, trans_a2, trans_b2, m2, n2, bs, &alpha, + static_cast(a2), lda2, static_cast(b2), ldb2, + &beta, static_cast(grad_weight->DataPtr()), ldc2)); + } + // d_bias = \sum_i(i=0, bs-1) d_output[i] + // TODO(dcj): use thrust::fill or reduce kernel do this + if (compute_grad_bias && bias) { + constexpr int BLOCK_SIZE = 256; + int threads_per_block = BLOCK_SIZE; + int num_blocks = out_features; + ReduceColumnsKernel<<>>( + static_cast(grad_output_promoted->DataPtr()), + static_cast(grad_bias->DataPtr()), out_features, bs); + } + }), + DataType::kFLOAT32) + DISPATCH_CASE(WRAP({ + if (compute_grad_input) { + CHECK(weight != nullptr) + << "compute_grad_input=true but weight is nullptr (selective save mismatch)"; + auto weight_promoted = weight_dtype == compute_dtype + ? weight + : std::make_shared(weight->To(compute_dtype)); + auto trans_a1 = transpose ? HIPBLAS_OP_N : HIPBLAS_OP_T; + auto lda1 = transpose ? in_features : out_features; + HIPBLAS_CHECK(hipblasGemmEx(handle, trans_a1, HIPBLAS_OP_N, in_features, bs, out_features, + &alpha, weight_promoted->DataPtr(), HIPBLAS_R_16B, lda1, + grad_output_promoted->DataPtr(), HIPBLAS_R_16B, out_features, + &beta, grad_input->DataPtr(), HIPBLAS_R_32F, in_features, + HIPBLAS_R_32F, HIPBLAS_GEMM_DEFAULT)); + } + if (compute_grad_weight) { + CHECK(input != nullptr) + << "compute_grad_weight=true but input is nullptr (selective save mismatch)"; + auto input_promoted = input_dtype == compute_dtype + ? input + : std::make_shared(input->To(compute_dtype)); + auto trans_a2 = HIPBLAS_OP_N; + auto trans_b2 = HIPBLAS_OP_T; + int m2 = transpose ? in_features : out_features; + int n2 = transpose ? out_features : in_features; + const void *a2 = transpose ? input_promoted->DataPtr() : grad_output_promoted->DataPtr(); + const void *b2 = transpose ? grad_output_promoted->DataPtr() : input_promoted->DataPtr(); + auto lda2 = transpose ? in_features : out_features; + auto ldb2 = transpose ? out_features : in_features; + auto ldc2 = transpose ? in_features : out_features; + HIPBLAS_CHECK(hipblasGemmEx(handle, trans_a2, trans_b2, m2, n2, bs, &alpha, a2, HIPBLAS_R_16B, + lda2, b2, HIPBLAS_R_16B, ldb2, &beta, grad_weight->DataPtr(), + HIPBLAS_R_32F, ldc2, HIPBLAS_R_32F, HIPBLAS_GEMM_DEFAULT)); + } + if (compute_grad_bias && bias) { + constexpr int BLOCK_SIZE = 256; + int threads_per_block = BLOCK_SIZE; + int num_blocks = out_features; + ReduceColumnsKernel<<>>( + static_cast(grad_output_promoted->DataPtr()), + static_cast(grad_bias->DataPtr()), out_features, bs); + } + }), + DataType::kBFLOAT16) + } + + return {grad_input, grad_weight, grad_bias}; +} +} // namespace infini_train::kernels::dcu + +#define REGISTER_HIP_LINEAR_KERNEL(kernel_name) \ + REGISTER_KERNEL(infini_train::Device::DeviceType::kDCU, kernel_name, infini_train::kernels::dcu::kernel_name) + +REGISTER_HIP_LINEAR_KERNEL(MatmulForward) +REGISTER_HIP_LINEAR_KERNEL(MatmulBackward) +REGISTER_HIP_LINEAR_KERNEL(LinearForward) +REGISTER_HIP_LINEAR_KERNEL(LinearBackward) + +#undef REGISTER_HIP_LINEAR_KERNEL diff --git a/infini_train/src/kernels/dcu/no_op.hip b/infini_train/src/kernels/dcu/no_op.hip new file mode 100644 index 00000000..392179a1 --- /dev/null +++ b/infini_train/src/kernels/dcu/no_op.hip @@ -0,0 +1,30 @@ +#include "glog/logging.h" + +#include "infini_train/include/dispatcher.h" +#include "infini_train/include/tensor.h" + +namespace infini_train::kernels::dcu { +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::dcu + +#define REGISTER_HIP_NO_OP_KERNEL(kernel_name) \ + REGISTER_KERNEL(infini_train::Device::DeviceType::kDCU, kernel_name, infini_train::kernels::dcu::kernel_name) + +REGISTER_HIP_NO_OP_KERNEL(NoOpForward) +REGISTER_HIP_NO_OP_KERNEL(NoOpBackward) + +#undef REGISTER_HIP_NO_OP_KERNEL diff --git a/infini_train/src/kernels/dcu/outer.hip b/infini_train/src/kernels/dcu/outer.hip new file mode 100644 index 00000000..7cba240b --- /dev/null +++ b/infini_train/src/kernels/dcu/outer.hip @@ -0,0 +1,166 @@ +#include +#include +#include + +#include + +#include "glog/logging.h" + +#include "infini_train/include/common/dcu/common_dcu.h" +#include "infini_train/include/core/runtime/device_guard.h" +#include "infini_train/include/dispatcher.h" +#include "infini_train/include/tensor.h" + +#include "infini_train/src/core/runtime/dcu/dcu_dispatch.h" +#include "infini_train/src/core/runtime/dcu/dcu_runtime_common.h" + +namespace infini_train::kernels::dcu { +std::shared_ptr OuterForward(const std::shared_ptr &input, const std::shared_ptr &other) { + /* + Computes outer product: output[i, j] = input[i] * other[j] + Equivalent to: input: [M, 1], other: [1, N] 鈫?output: [M, N] + */ + + const auto &in_dims = input->Dims(); + const auto &ot_dims = other->Dims(); + // TODO(zbl): support batched outer? + CHECK_EQ(in_dims.size(), 1); + CHECK_EQ(ot_dims.size(), 1); + + const int64_t M = in_dims[0]; + const int64_t N = ot_dims[0]; + + auto output = std::make_shared(std::vector{M, N}, input->Dtype(), input->GetDevice()); + + auto device = input->GetDevice(); + // reinterpret input: [M] as column vector [M, 1] + // reinterpret other: [N] as row vector [1, N] + // output[M, N] = input[M, 1] * other.T[1, N] + // output.T[N, M] = other[N, 1] * input.T[1, M] + float alpha = 1.0f; + float beta = 0.0f; + hipblasHandle_t handle = dynamic_cast( + infini_train::core::GetDeviceGuardImpl(device.type())->GetBlasHandle(device)) + ->hipblas_handle(); + + switch (input->Dtype()) { + DISPATCH_CASE(WRAP({ + HIPBLAS_CHECK(hipblasSgemm(handle, HIPBLAS_OP_N, HIPBLAS_OP_N, N, M, 1, &alpha, + static_cast(other->DataPtr()), N, + static_cast(input->DataPtr()), 1, &beta, + static_cast(output->DataPtr()), N)); + }), + DataType::kFLOAT32) + DISPATCH_CASE(WRAP({ + HIPBLAS_CHECK(hipblasGemmEx(handle, HIPBLAS_OP_N, HIPBLAS_OP_N, N, M, 1, &alpha, other->DataPtr(), + HIPBLAS_R_16B, N, input->DataPtr(), HIPBLAS_R_16B, 1, &beta, + output->DataPtr(), HIPBLAS_R_16B, N, HIPBLAS_R_32F, HIPBLAS_GEMM_DEFAULT)); + }), + DataType::kBFLOAT16) + } + + return output; +} + +std::tuple, std::shared_ptr> OuterBackward(const std::shared_ptr &input, + const std::shared_ptr &other, + const std::shared_ptr &grad_output) { + /* + grad_input: [M] = grad_output: [M, N] 脳 other: [N] + grad_other: [N] = grad_output.T: [N, M] 脳 input: [M] + */ + const int64_t M = input->Dims()[0]; + const int64_t N = other->Dims()[0]; + // TODO(zbl): support batched outer? + CHECK_EQ(grad_output->Dims().size(), 2); + CHECK_EQ(grad_output->Dims()[0], M); + CHECK_EQ(grad_output->Dims()[1], N); + + auto input_dtype = input->Dtype(); + auto other_dtype = other->Dtype(); + auto grad_output_dtype = grad_output->Dtype(); + + // Compute dtype determined by saved tensors (forward compute dtype), not grad_output + DataType promoted_type = PromoteDataTypes(input_dtype, other_dtype); + + auto input_promoted = input_dtype == promoted_type ? input : std::make_shared(input->To(promoted_type)); + auto other_promoted = other_dtype == promoted_type ? other : std::make_shared(other->To(promoted_type)); + auto grad_output_promoted + = grad_output_dtype == promoted_type ? grad_output : std::make_shared(grad_output->To(promoted_type)); + + // For bf16 compute, output in fp32 to preserve accumulation precision (matches PyTorch behavior) + auto output_dtype = (promoted_type == DataType::kBFLOAT16) ? DataType::kFLOAT32 : promoted_type; + auto grad_input = std::make_shared(std::vector{M}, output_dtype, grad_output->GetDevice()); + auto grad_other = std::make_shared(std::vector{N}, output_dtype, grad_output->GetDevice()); + + core::dcu::DispatchDcuFunc( + promoted_type, + [=]() { + grad_input->Fill(0.0); + grad_other->Fill(0.0); + }, + "HIP OuterBackward"); + + auto device = input->GetDevice(); + float alpha = 1.0f; + float beta = 0.0f; + hipblasHandle_t handle = dynamic_cast( + infini_train::core::GetDeviceGuardImpl(device.type())->GetBlasHandle(device)) + ->hipblas_handle(); + + switch (promoted_type) { + DISPATCH_CASE(WRAP({ + // grad_input[M, 1] = grad_output[M, N] 脳 other[N, 1] + // y = grad_input[M] + // A = grad_output.T[N, M] + // x = other[N] + HIPBLAS_CHECK(hipblasSgemv(handle, HIPBLAS_OP_T, N, M, &alpha, + static_cast(grad_output_promoted->DataPtr()), N, + static_cast(other_promoted->DataPtr()), 1, &beta, + static_cast(grad_input->DataPtr()), 1)); + + // grad_other[N, 1] = grad_output.T[N, M] 脳 input[M, 1] + // y = grad_other[N] + // A = grad_output.T[N, M] + // x = input[M] + HIPBLAS_CHECK(hipblasSgemv(handle, HIPBLAS_OP_N, N, M, &alpha, + static_cast(grad_output_promoted->DataPtr()), N, + static_cast(input_promoted->DataPtr()), 1, &beta, + static_cast(grad_other->DataPtr()), 1)); + }), + DataType::kFLOAT32) + DISPATCH_CASE( + // hipblasgemv does not support bf16, use hipblasGemmEx to workaround + WRAP({ + // grad_input[M, 1] = grad_output[M, N] 脳 other[N, 1] + // grad_input.T[1, M] = other.T[1, N] 脳 grad_output.T[N, M] + // C = grad_input.T[1, M] + // A = other.T[1, N] + // B = grad_output.T[N, M] + HIPBLAS_CHECK(hipblasGemmEx(handle, HIPBLAS_OP_N, HIPBLAS_OP_N, 1, M, N, &alpha, other_promoted->DataPtr(), + HIPBLAS_R_16B, 1, grad_output_promoted->DataPtr(), HIPBLAS_R_16B, N, &beta, + grad_input->DataPtr(), HIPBLAS_R_32F, 1, HIPBLAS_R_32F, HIPBLAS_GEMM_DEFAULT)); + // grad_other[N, 1] = grad_output.T[N, M] 脳 input[M, 1] + // grad_other.T[1, N] = input.T[1, M] 脳 grad_output[M, N] + // C = grad_other.T[1, N] + // A = input.T[1, M] + // B = grad_output.T[N, M] + HIPBLAS_CHECK(hipblasGemmEx(handle, HIPBLAS_OP_N, HIPBLAS_OP_T, 1, N, M, &alpha, input_promoted->DataPtr(), + HIPBLAS_R_16B, 1, grad_output_promoted->DataPtr(), HIPBLAS_R_16B, N, &beta, + grad_other->DataPtr(), HIPBLAS_R_32F, 1, HIPBLAS_R_32F, HIPBLAS_GEMM_DEFAULT)); + }), + DataType::kBFLOAT16) + } + + return {grad_input, grad_other}; +} + +} // namespace infini_train::kernels::dcu + +#define REGISTER_HIP_OUTER_KERNEL(kernel_name) \ + REGISTER_KERNEL(infini_train::Device::DeviceType::kDCU, kernel_name, infini_train::kernels::dcu::kernel_name) + +REGISTER_HIP_OUTER_KERNEL(OuterForward) +REGISTER_HIP_OUTER_KERNEL(OuterBackward) + +#undef REGISTER_HIP_OUTER_KERNEL diff --git a/infini_train/src/kernels/dcu/reduction.hip b/infini_train/src/kernels/dcu/reduction.hip new file mode 100644 index 00000000..c37efce0 --- /dev/null +++ b/infini_train/src/kernels/dcu/reduction.hip @@ -0,0 +1,244 @@ +#include + +#include "infini_train/include/common/dcu/common_dcu.h" +#include "infini_train/include/common/dcu/cub_compat.cuh" +#include "infini_train/include/common/dcu/kernel_helper.cuh" +#include "infini_train/include/core/runtime/device_guard.h" +#include "infini_train/include/dispatcher.h" +#include "infini_train/include/tensor.h" + +#include "infini_train/src/core/runtime/dcu/dcu_dispatch.h" +#include "infini_train/src/core/runtime/dcu/dcu_runtime_common.h" + +namespace infini_train::kernels::dcu { +namespace { +__host__ __device__ constexpr float kInfinity = std::numeric_limits::infinity(); +} // namespace + +namespace { +// Reduction operators +template struct CubOp; + +template struct CubOp { + __device__ static T Init() { return common::dcu::Cast(0); } + __device__ static T Reduce(T a, T b) { return common::dcu::Add(a, b); } + __device__ static CubSumOp Op() { return CubSumOp(); } +}; + +template struct CubOp { + __device__ static T Init() { return common::dcu::Cast(-kInfinity); } + __device__ static T Reduce(T a, T b) { return common::dcu::Max(a, b); } + __device__ static CubMaxOp Op() { return CubMaxOp(); } +}; + +template struct CubOp { + __device__ static T Init() { return common::dcu::Cast(kInfinity); } + __device__ static T Reduce(T a, T b) { return common::dcu::Min(a, b); } + __device__ static CubMinOp Op() { return CubMinOp(); } +}; + +// Finalization strategies +template struct MeanFinalize { + __device__ __forceinline__ T operator()(T sum, int64_t count) const { + return common::dcu::Div(sum, common::dcu::Cast(count)); + } +}; + +template struct IdentityFinalize { + __device__ __forceinline__ T operator()(T val, int64_t) const { return val; } +}; + +// Generic reduction kernel +template +__global__ void GenericReduceKernel(const T *input, T *output, int64_t N, int64_t H, int64_t W, + FinalizeOp finalize_op) { + using BlockReduce = hipcub::BlockReduce; + __shared__ typename BlockReduce::TempStorage temp_storage; + + int idx = blockIdx.x; + if (idx >= N * W) { + return; + } + + int n = idx / W; + int w = idx % W; + + T acc = CubOp::Init(); + for (int64_t h = threadIdx.x; h < H; h += blockDim.x) { + int input_idx = (n * H + h) * W + w; + acc = CubOp::Reduce(acc, input[input_idx]); + } + + T reduced = BlockReduce(temp_storage).Reduce(acc, CubOp::Op()); + + if (threadIdx.x == 0) { + output[idx] = finalize_op(reduced, H); + } +} + +// Unified backward kernel for Mean, Sum, Max, and Min +template +__global__ void GenericReduceBackwardKernel(T *grad_input, const T *grad_output, const T *input, const T *reduced, + int64_t N, int64_t H, int64_t W, bool is_mean, bool is_masked) { + int idx = blockIdx.x * blockDim.x + threadIdx.x; + if (idx >= N * H * W) { + return; + } + + int n = idx / (H * W); + int hw = idx % (H * W); + int w = hw % W; + + int reduced_idx = n * W + w; + + if (is_masked) { + T selected = reduced[reduced_idx]; + T value = input[idx]; + grad_input[idx] = (value == selected) ? grad_output[reduced_idx] : T(0); + } else { + grad_input[idx] = grad_output[reduced_idx]; + if (is_mean) { + T H_casted; + // TODO(lzm): directly use Cast when (half and hip_bfloat16) <-> (integral types) is supported + if constexpr (std::is_same_v || std::is_same_v) { + H_casted = common::dcu::Cast(static_cast(H)); + } else { + H_casted = common::dcu::Cast(H); + } + grad_input[idx] /= H_casted; + } + } +} +} // namespace + +// Common forward implementation for reduce ops +template class FinalizeOp> +std::shared_ptr ReduceOpForward(const std::shared_ptr &input, const int64_t dim, const bool keep_dim) { + const auto &input_dims = input->Dims(); + int64_t actual_dim = dim < 0 ? dim + input_dims.size() : dim; + CHECK_GE(actual_dim, 0); + CHECK_LT(actual_dim, input_dims.size()); + + std::vector output_dims = input_dims; + if (keep_dim) { + output_dims[actual_dim] = 1; + } else { + output_dims.erase(output_dims.begin() + actual_dim); + } + + auto dtype = input->Dtype(); + auto output = std::make_shared(output_dims, dtype, input->GetDevice()); + + int64_t N = std::accumulate(input_dims.begin(), input_dims.begin() + actual_dim, 1, std::multiplies()); + int64_t H = input_dims[actual_dim]; + int64_t W = std::accumulate(input_dims.begin() + actual_dim + 1, input_dims.end(), 1, std::multiplies()); + + constexpr int BLOCK_SIZE = 256; + int threads_per_block = BLOCK_SIZE; + int num_blocks = N * W; + + auto device = input->GetDevice(); + const auto &hip_stream = dynamic_cast( + infini_train::core::GetDeviceGuardImpl(device.type())->GetStream(device)) + ->hip_stream(); + + core::dcu::DispatchDcuFunc( + dtype, + [=]() { + GenericReduceKernel, BLOCK_SIZE> + <<>>(static_cast(input->DataPtr()), + static_cast(output->DataPtr()), N, H, W, + FinalizeOp{}); + }, + "HIP ReductionForward"); + return output; +} + +// Common backward implementation for reduce ops +std::shared_ptr ReduceOpBackward(const std::shared_ptr &grad_output, + const std::shared_ptr &input, const std::shared_ptr &reduced, + const std::vector &input_dims, const int64_t dim, bool keep_dim, + bool is_mean, bool is_masked) { + int64_t actual_dim = dim < 0 ? dim + input_dims.size() : dim; + CHECK_GE(actual_dim, 0); + CHECK_LT(actual_dim, input_dims.size()); + + auto dtype = grad_output->Dtype(); + auto grad_input = std::make_shared(input_dims, dtype, grad_output->GetDevice()); + + int64_t N = std::accumulate(input_dims.begin(), input_dims.begin() + actual_dim, 1, std::multiplies()); + int64_t H = input_dims[actual_dim]; + int64_t W = std::accumulate(input_dims.begin() + actual_dim + 1, input_dims.end(), 1, std::multiplies()); + + int threads_per_block = 256; + int num_blocks = (N * H * W + threads_per_block - 1) / threads_per_block; + + auto device = grad_output->GetDevice(); + const auto &hip_stream = dynamic_cast( + infini_train::core::GetDeviceGuardImpl(device.type())->GetStream(device)) + ->hip_stream(); + + core::dcu::DispatchDcuFunc( + dtype, + [=]() { + grad_input->Fill(0.0); + GenericReduceBackwardKernel<<>>( + static_cast(grad_input->DataPtr()), static_cast(grad_output->DataPtr()), + input ? static_cast(input->DataPtr()) : nullptr, + reduced ? static_cast(reduced->DataPtr()) : nullptr, N, H, W, is_mean, is_masked); + }, + "HIP ReductionBackward"); + return grad_input; +} + +std::shared_ptr MeanForward(const std::shared_ptr &input, const int64_t dim, const bool keep_dim) { + return ReduceOpForward(input, dim, keep_dim); +} + +std::shared_ptr SumForward(const std::shared_ptr &input, const int64_t dim, const bool keep_dim) { + return ReduceOpForward(input, dim, keep_dim); +} + +std::shared_ptr MaxForward(const std::shared_ptr &input, const int64_t dim, const bool keep_dim) { + return ReduceOpForward(input, dim, keep_dim); +} + +std::shared_ptr MinForward(const std::shared_ptr &input, const int64_t dim, const bool keep_dim) { + return ReduceOpForward(input, dim, keep_dim); +} + +std::shared_ptr MeanBackward(const std::shared_ptr &grad_output, const std::vector &input_dims, + const int64_t dim, bool keep_dim) { + return ReduceOpBackward(grad_output, nullptr, nullptr, input_dims, dim, keep_dim, true, false); +} + +std::shared_ptr SumBackward(const std::shared_ptr &grad_output, const std::vector &input_dims, + const int64_t dim, bool keep_dim) { + return ReduceOpBackward(grad_output, nullptr, nullptr, input_dims, dim, keep_dim, false, false); +} + +std::shared_ptr MaxBackward(const std::shared_ptr &grad_output, const std::shared_ptr &input, + const std::shared_ptr &reduced, const int64_t dim, bool keep_dim) { + return ReduceOpBackward(grad_output, input, reduced, input->Dims(), dim, keep_dim, false, true); +} + +std::shared_ptr MinBackward(const std::shared_ptr &grad_output, const std::shared_ptr &input, + const std::shared_ptr &reduced, const int64_t dim, bool keep_dim) { + return ReduceOpBackward(grad_output, input, reduced, input->Dims(), dim, keep_dim, false, true); +} + +} // namespace infini_train::kernels::dcu + +#define REGISTER_HIP_REDUCTION_KERNEL(kernel_name) \ + REGISTER_KERNEL(infini_train::Device::DeviceType::kDCU, kernel_name, infini_train::kernels::dcu::kernel_name) + +REGISTER_HIP_REDUCTION_KERNEL(MeanForward) +REGISTER_HIP_REDUCTION_KERNEL(SumForward) +REGISTER_HIP_REDUCTION_KERNEL(MaxForward) +REGISTER_HIP_REDUCTION_KERNEL(MinForward) +REGISTER_HIP_REDUCTION_KERNEL(MeanBackward) +REGISTER_HIP_REDUCTION_KERNEL(SumBackward) +REGISTER_HIP_REDUCTION_KERNEL(MaxBackward) +REGISTER_HIP_REDUCTION_KERNEL(MinBackward) + +#undef REGISTER_HIP_REDUCTION_KERNEL diff --git a/infini_train/src/kernels/dcu/slice.hip b/infini_train/src/kernels/dcu/slice.hip new file mode 100644 index 00000000..17cf0144 --- /dev/null +++ b/infini_train/src/kernels/dcu/slice.hip @@ -0,0 +1,209 @@ +#include +#include + +#include "glog/logging.h" + +#include "infini_train/include/common/dcu/common_dcu.h" +#include "infini_train/include/core/runtime/device_guard.h" +#include "infini_train/include/dispatcher.h" +#include "infini_train/include/tensor.h" + +#include "infini_train/src/core/runtime/dcu/dcu_dispatch.h" +#include "infini_train/src/core/runtime/dcu/dcu_runtime_common.h" + +namespace infini_train::kernels::dcu { + +template +__global__ void SliceForwardKernel(const T *input, T *output, const int64_t *new_dims, const int64_t *starts, + const int64_t *steps, const int64_t *in_strides, const int64_t *out_strides, + int num_dims, int64_t total_elements) { + int64_t out_idx = blockIdx.x * blockDim.x + threadIdx.x; + if (out_idx >= total_elements) { + return; + } + + int64_t in_index = 0; + for (int i = 0; i < num_dims; ++i) { + int64_t idx = (out_idx / out_strides[i]) % new_dims[i]; + in_index += (starts[i] + idx * steps[i]) * in_strides[i]; + } + + output[out_idx] = input[in_index]; +} + +std::shared_ptr SliceForward(const std::shared_ptr &input, const std::vector &starts, + const std::vector &ends, const std::vector &steps) { + CHECK_EQ(starts.size(), ends.size()); + CHECK_EQ(starts.size(), steps.size()); + auto &dims = input->Dims(); + CHECK_EQ(starts.size(), dims.size()); + const int64_t num_dims = dims.size(); + + std::vector new_dims; + for (int i = 0; i < starts.size(); ++i) { + CHECK_LE(starts[i], ends[i]); + CHECK_LE(0, steps[i]); + new_dims.push_back((ends[i] - starts[i] + steps[i] - 1) / steps[i]); + } + + auto dtype = input->Dtype(); + auto new_tensor = std::make_shared(new_dims, dtype, input->GetDevice()); + // NOTE(zbl): must initialize with 0 + new_tensor->Fill(0.0); + + std::vector src_strides(dims.size(), 0), dst_strides(new_dims.size(), 0); + int64_t stride = 1; + for (int i = dims.size() - 1; i >= 0; --i) { + src_strides[i] = stride; + stride *= dims[i]; + } + + stride = 1; + for (int i = new_dims.size() - 1; i >= 0; --i) { + dst_strides[i] = stride; + stride *= new_dims[i]; + } + + int64_t total_elements = stride; + + int64_t *new_dims_dev, *starts_dev, *steps_dev, *input_strides_dev, *output_strides_dev; + + auto device = input->GetDevice(); + const auto &stream = dynamic_cast( + infini_train::core::GetDeviceGuardImpl(device.type())->GetStream(device)) + ->hip_stream(); + + hipMallocAsync(&new_dims_dev, + (ends.size() + starts.size() + steps.size() + dims.size() + new_dims.size()) * sizeof(int64_t), + stream); + starts_dev = new_dims_dev + ends.size(); + steps_dev = starts_dev + starts.size(); + input_strides_dev = steps_dev + steps.size(); + output_strides_dev = input_strides_dev + dims.size(); + + hipMemcpyAsync(new_dims_dev, new_dims.data(), ends.size() * sizeof(int64_t), hipMemcpyHostToDevice, stream); + hipMemcpyAsync(starts_dev, starts.data(), starts.size() * sizeof(int64_t), hipMemcpyHostToDevice, stream); + hipMemcpyAsync(steps_dev, steps.data(), steps.size() * sizeof(int64_t), hipMemcpyHostToDevice, stream); + hipMemcpyAsync(input_strides_dev, src_strides.data(), dims.size() * sizeof(int64_t), hipMemcpyHostToDevice, + stream); + hipMemcpyAsync(output_strides_dev, dst_strides.data(), new_dims.size() * sizeof(int64_t), hipMemcpyHostToDevice, + stream); + + int threads_per_block = 256; + int num_blocks = (total_elements + threads_per_block - 1) / threads_per_block; + + core::dcu::DispatchDcuFunc( + dtype, + [=]() { + SliceForwardKernel<<>>( + static_cast(input->DataPtr()), static_cast(new_tensor->DataPtr()), new_dims_dev, + starts_dev, steps_dev, input_strides_dev, output_strides_dev, num_dims, total_elements); + }, + "HIP SliceForward"); + + HIP_CHECK(hipFree(new_dims_dev)); + + return new_tensor; +} + +template +__global__ void SliceBackwardKernel(const T *grad_output, T *grad_input, const int64_t *new_dims, const int64_t *starts, + const int64_t *steps, const int64_t *in_strides, const int64_t *out_strides, + int num_dims, int64_t total_elements) { + int64_t out_idx = blockIdx.x * blockDim.x + threadIdx.x; + if (out_idx >= total_elements) { + return; + } + + int64_t in_index = 0; + for (int i = 0; i < num_dims; ++i) { + int64_t idx = (out_idx / out_strides[i]) % new_dims[i]; + in_index += (starts[i] + idx * steps[i]) * in_strides[i]; + } + grad_input[in_index] = grad_output[out_idx]; +} + +std::shared_ptr SliceBackward(const std::shared_ptr &grad_output, const std::shared_ptr &input, + const std::vector &starts, const std::vector &ends, + const std::vector &steps) { + CHECK_EQ(starts.size(), ends.size()); + CHECK_EQ(starts.size(), steps.size()); + auto &dims = input->Dims(); + CHECK_EQ(starts.size(), dims.size()); + const int64_t num_dims = dims.size(); + + std::vector new_dims; + for (int i = 0; i < starts.size(); ++i) { + CHECK_LE(starts[i], ends[i]); + CHECK_LE(0, steps[i]); + new_dims.push_back((ends[i] - starts[i] + steps[i] - 1) / steps[i]); + } + + auto grad_output_dtype = grad_output->Dtype(); + auto grad_input = std::make_shared(input->Dims(), grad_output_dtype, grad_output->GetDevice()); + grad_input->Fill(0.0); + + std::vector src_strides(dims.size()); + int64_t stride = 1; + for (int i = src_strides.size() - 1; i >= 0; --i) { + src_strides[i] = stride; + stride *= dims[i]; + } + + std::vector dst_strides(new_dims.size()); + stride = 1; + for (int i = dst_strides.size() - 1; i >= 0; --i) { + dst_strides[i] = stride; + stride *= new_dims[i]; + } + + int64_t total_elements = stride; + + int dims_size = dims.size(); + int64_t *new_dims_dev, *starts_dev, *steps_dev, *input_strides_dev, *output_strides_dev; + + auto device = input->GetDevice(); + const auto &stream = dynamic_cast( + infini_train::core::GetDeviceGuardImpl(device.type())->GetStream(device)) + ->hip_stream(); + hipMallocAsync(&new_dims_dev, + (ends.size() + starts.size() + steps.size() + dims.size() + new_dims.size()) * sizeof(int64_t), + stream); + starts_dev = new_dims_dev + ends.size(); + steps_dev = starts_dev + starts.size(); + input_strides_dev = steps_dev + steps.size(); + output_strides_dev = input_strides_dev + dims.size(); + + hipMemcpyAsync(new_dims_dev, new_dims.data(), ends.size() * sizeof(int64_t), hipMemcpyHostToDevice, stream); + hipMemcpyAsync(starts_dev, starts.data(), starts.size() * sizeof(int64_t), hipMemcpyHostToDevice, stream); + hipMemcpyAsync(steps_dev, steps.data(), steps.size() * sizeof(int64_t), hipMemcpyHostToDevice, stream); + hipMemcpyAsync(input_strides_dev, src_strides.data(), dims.size() * sizeof(int64_t), hipMemcpyHostToDevice, + stream); + hipMemcpyAsync(output_strides_dev, dst_strides.data(), new_dims.size() * sizeof(int64_t), hipMemcpyHostToDevice, + stream); + + int threads_per_block = 256; + int num_blocks = (total_elements + threads_per_block - 1) / threads_per_block; + + core::dcu::DispatchDcuFunc( + grad_output_dtype, + [=]() { + SliceBackwardKernel<<>>( + static_cast(grad_output->DataPtr()), static_cast(grad_input->DataPtr()), new_dims_dev, + starts_dev, steps_dev, input_strides_dev, output_strides_dev, num_dims, total_elements); + }, + "HIP SliceBackward"); + + HIP_CHECK(hipFree(new_dims_dev)); + + return grad_input; +} +} // namespace infini_train::kernels::dcu + +#define REGISTER_HIP_SLICE_KERNEL(kernel_name) \ + REGISTER_KERNEL(infini_train::Device::DeviceType::kDCU, kernel_name, infini_train::kernels::dcu::kernel_name) + +REGISTER_HIP_SLICE_KERNEL(SliceForward) +REGISTER_HIP_SLICE_KERNEL(SliceBackward) + +#undef REGISTER_HIP_SLICE_KERNEL diff --git a/infini_train/src/kernels/dcu/softmax.hip b/infini_train/src/kernels/dcu/softmax.hip new file mode 100644 index 00000000..b09d44d3 --- /dev/null +++ b/infini_train/src/kernels/dcu/softmax.hip @@ -0,0 +1,222 @@ +#include +#include + +#include + +#include "glog/logging.h" + +#include "infini_train/include/common/dcu/common_dcu.h" +#include "infini_train/include/common/dcu/cub_compat.cuh" +#include "infini_train/include/common/dcu/kernel_helper.cuh" +#include "infini_train/include/core/runtime/device_guard.h" +#include "infini_train/include/dispatcher.h" +#include "infini_train/include/tensor.h" + +#include "infini_train/src/core/runtime/dcu/dcu_dispatch.h" +#include "infini_train/src/core/runtime/dcu/dcu_runtime_common.h" + +namespace infini_train::kernels::dcu { +template +__global__ void SoftmaxForwardKernel(T *output, const T *input, int64_t outer_size, int64_t axis_size, + int64_t inner_size) { + using BlockReduce = hipcub::BlockReduce; + + __shared__ typename BlockReduce::TempStorage temp_storage_max; + __shared__ typename BlockReduce::TempStorage temp_storage_sum; + __shared__ float row_max; + __shared__ float row_sum; + + const int64_t group = blockIdx.x; // row of the grid + const int64_t inner_idx = blockIdx.y; // column of the grid + const int tid = threadIdx.x; + + // calculate the maximum for each group + float thread_max = -INFINITY; + for (int64_t axis = tid; axis < axis_size; axis += BLOCK_SIZE) { + int64_t idx = (group * axis_size + axis) * inner_size + inner_idx; + thread_max = max(thread_max, common::dcu::Cast(input[idx])); + } + float block_max = BlockReduce(temp_storage_max).Reduce(thread_max, CubMaxOp()); + + if (tid == 0) { + row_max = block_max; + } + __syncthreads(); + + // calculate the sum of exponents + float thread_sum = 0; + for (int64_t axis = tid; axis < axis_size; axis += BLOCK_SIZE) { + int64_t idx = (group * axis_size + axis) * inner_size + inner_idx; + float exp_val = exp(common::dcu::Cast(input[idx]) - row_max); + output[idx] = common::dcu::Cast(exp_val); + thread_sum += exp_val; + } + float block_sum = BlockReduce(temp_storage_sum).Sum(thread_sum); + + if (tid == 0) { + row_sum = block_sum; + } + __syncthreads(); + + // normalize + for (int64_t axis = tid; axis < axis_size; axis += BLOCK_SIZE) { + int64_t idx = (group * axis_size + axis) * inner_size + inner_idx; + output[idx] = common::dcu::Cast(common::dcu::Cast(output[idx]) / row_sum); + } +} + +template +void LaunchForward(const std::shared_ptr &output, const std::shared_ptr &input, int64_t dim) { + const auto &input_dims = input->Dims(); + int64_t outer_size = 1; + int64_t axis_size = input_dims[dim]; + int64_t inner_size = 1; + + for (int i = 0; i < dim; ++i) { outer_size *= input_dims[i]; }; + for (int i = dim + 1; i < input_dims.size(); ++i) { inner_size *= input_dims[i]; }; + if (axis_size == 0) { + LOG_LOC(INFO, "HIP softmax forward: 'input_dims[dim] == 0'"); + return; + } + if (outer_size == 0) { + return; + } + + T *output_ptr = static_cast(output->DataPtr()); + const T *input_ptr = static_cast(input->DataPtr()); + + if (BLOCK_SIZE > 1024) { + LOG_LOC(FATAL, "HIP softmax forward: 'BLOCK_SIZE used is larger than the max number of thread per block'"); + } + dim3 block_dims(BLOCK_SIZE); + dim3 grid_dims(outer_size, inner_size); + + auto device = output->GetDevice(); + const auto &hip_stream = dynamic_cast( + infini_train::core::GetDeviceGuardImpl(device.type())->GetStream(device)) + ->hip_stream(); + SoftmaxForwardKernel + <<>>(output_ptr, input_ptr, outer_size, axis_size, inner_size); +} + +std::shared_ptr SoftmaxForward(const std::shared_ptr &input, int64_t dim) { + auto dtype = input->Dtype(); + const auto &input_dims = input->Dims(); + dim = dim < 0 ? dim + input_dims.size() : dim; + CHECK(dim >= 0 && dim < input_dims.size()); + auto output = std::make_shared(input_dims, dtype, input->GetDevice()); + + switch (dtype) { + DISPATCH_CASE(WRAP(LaunchForward<256, float>(output, input, dim);), DataType::kFLOAT32) + DISPATCH_CASE(WRAP(LaunchForward<256, hip_bfloat16>(output, input, dim);), DataType::kBFLOAT16) + default: + LOG_LOC(FATAL, "HIP softmax forward: 'Unsupported data type'"); + } + return output; +} + +template +__global__ void SoftmaxBackwardKernel(T *grad_input, const T *grad_output, const T *output, int64_t outer_size, + int64_t axis_size, int64_t inner_size) { + using BlockReduce = hipcub::BlockReduce; + + __shared__ typename BlockReduce::TempStorage temp_storage_sum; + __shared__ float row_sum; + + const int64_t group = blockIdx.x; + const int64_t inner_idx = blockIdx.y; + const int tid = threadIdx.x; + + // calculate the sum of the dot product of gradients + float thread_sum = 0; + for (int64_t axis = tid; axis < axis_size; axis += BLOCK_SIZE) { + const int64_t idx = (group * axis_size + axis) * inner_size + inner_idx; + thread_sum += common::dcu::Cast(grad_output[idx] * output[idx]); + } + float block_sum = BlockReduce(temp_storage_sum).Sum(thread_sum); + + if (tid == 0) { + row_sum = block_sum; + } + __syncthreads(); + + // update the input gradient + for (int64_t axis = tid; axis < axis_size; axis += BLOCK_SIZE) { + const int64_t idx = (group * axis_size + axis) * inner_size + inner_idx; + grad_input[idx] = output[idx] * (grad_output[idx] - common::dcu::Cast(row_sum)); + } +} + +template +void LaunchBackward(const std::shared_ptr &grad_input, const std::shared_ptr &grad_output, + const std::shared_ptr &output, int64_t dim) { + const auto &output_dims = output->Dims(); + int64_t outer_size = 1; + int64_t axis_size = output_dims[dim]; + int64_t inner_size = 1; + + for (int i = 0; i < dim; ++i) { outer_size *= output_dims[i]; }; + for (int i = dim + 1; i < output_dims.size(); ++i) { inner_size *= output_dims[i]; }; + if (axis_size == 0) { + LOG_LOC(INFO, "HIP softmax backward: 'output_dims[dim] == 0'"); + return; + } + if (outer_size == 0) { + return; + } + + T *grad_input_ptr = static_cast(grad_input->DataPtr()); + const T *grad_output_ptr = static_cast(grad_output->DataPtr()); + const T *output_ptr = static_cast(output->DataPtr()); + + if (BLOCK_SIZE > 1024) { + LOG_LOC(FATAL, "HIP softmax backward: 'BLOCK_SIZE used is larger than the max number of thread per block'"); + } + dim3 block(BLOCK_SIZE); + dim3 grid(outer_size, inner_size); + + auto device = output->GetDevice(); + const auto &hip_stream = dynamic_cast( + infini_train::core::GetDeviceGuardImpl(device.type())->GetStream(device)) + ->hip_stream(); + SoftmaxBackwardKernel<<>>(grad_input_ptr, grad_output_ptr, output_ptr, + outer_size, axis_size, inner_size); +} + +std::shared_ptr SoftmaxBackward(const std::shared_ptr &grad_output, + const std::shared_ptr &output, int64_t dim) { + auto grad_output_dtype = grad_output->Dtype(); + auto output_dtype = output->Dtype(); + DataType promoted_type = PromoteDataTypes(grad_output_dtype, output_dtype); + + auto grad_output_promoted + = grad_output_dtype == promoted_type ? grad_output : std::make_shared(grad_output->To(promoted_type)); + auto output_promoted = output_dtype == promoted_type ? output : std::make_shared(output->To(promoted_type)); + + const auto &output_dims = output->Dims(); + dim = dim < 0 ? dim + output->Dims().size() : dim; + CHECK(dim >= 0 && dim < output->Dims().size()); + + auto grad_input = std::make_shared(output_dims, promoted_type, output->GetDevice()); + grad_input->Fill(0.0); + + switch (promoted_type) { + DISPATCH_CASE(WRAP(LaunchBackward<256, float>(grad_input, grad_output_promoted, output_promoted, dim);), + DataType::kFLOAT32) + DISPATCH_CASE(WRAP(LaunchBackward<256, hip_bfloat16>(grad_input, grad_output_promoted, output_promoted, dim);), + DataType::kBFLOAT16) + default: + LOG_LOC(FATAL, "HIP softmax backward: 'Unsupported data type'"); + } + + return grad_input; +} +} // namespace infini_train::kernels::dcu + +#define REGISTER_HIP_SOFTMAX_KERNEL(kernel_name) \ + REGISTER_KERNEL(infini_train::Device::DeviceType::kDCU, kernel_name, infini_train::kernels::dcu::kernel_name) + +REGISTER_HIP_SOFTMAX_KERNEL(SoftmaxForward) +REGISTER_HIP_SOFTMAX_KERNEL(SoftmaxBackward) + +#undef REGISTER_HIP_SOFTMAX_KERNEL diff --git a/infini_train/src/kernels/dcu/split.hip b/infini_train/src/kernels/dcu/split.hip new file mode 100644 index 00000000..90c9739e --- /dev/null +++ b/infini_train/src/kernels/dcu/split.hip @@ -0,0 +1,182 @@ +#include +#include +#include + +#include "infini_train/include/common/dcu/common_dcu.h" +#include "infini_train/include/core/runtime/device_guard.h" +#include "infini_train/include/dispatcher.h" +#include "infini_train/include/tensor.h" + +#include "infini_train/src/core/runtime/dcu/dcu_dispatch.h" +#include "infini_train/src/core/runtime/dcu/dcu_runtime_common.h" + +namespace infini_train::kernels::dcu { +template +__global__ void SplitForwardKernel(const T *input, T *output, int64_t N, int64_t H_in, int64_t H_out, int64_t W, + int64_t start_idx) { + int idx = blockIdx.x * blockDim.x + threadIdx.x; + int total = N * H_out * W; + + if (idx < total) { + int w = idx % W; + int h = (idx / W) % H_out; + int n = idx / (H_out * W); + + int input_h = h + start_idx; + int input_idx = n * H_in * W + input_h * W + w; + int output_idx = n * H_out * W + h * W + w; + + output[output_idx] = input[input_idx]; + } +} + +std::vector> SplitForward(const std::shared_ptr &input, int64_t split_size, int dim) { + CHECK_GT(split_size, 0); + CHECK_GE(dim, 0) << "Currently we do not support negative dimension"; + const auto &input_dims = input->Dims(); + CHECK_LT(dim, input_dims.size()); + + std::vector> outputs; + auto dtype = input->Dtype(); + + const int64_t N = std::accumulate(input_dims.begin(), input_dims.begin() + dim, 1, std::multiplies()); + const int64_t W = std::accumulate(input_dims.begin() + dim + 1, input_dims.end(), 1, std::multiplies()); + const int64_t H_in = input_dims[dim]; + + for (int64_t start = 0; start < H_in; start += split_size) { + auto output_dims = input_dims; + const int64_t H_out = std::min(split_size, H_in - start); + output_dims[dim] = H_out; + + auto output = std::make_shared(output_dims, dtype, input->GetDevice()); + + int64_t total = N * H_out * W; + int threads_per_block = 256; + int num_blocks = (total + threads_per_block - 1) / threads_per_block; + + auto device = input->GetDevice(); + const auto &hip_stream = dynamic_cast( + infini_train::core::GetDeviceGuardImpl(device.type())->GetStream(device)) + ->hip_stream(); + + core::dcu::DispatchDcuFunc( + dtype, + [=]() { + SplitForwardKernel<<>>( + static_cast(input->DataPtr()), static_cast(output->DataPtr()), N, H_in, H_out, W, + start); + }, + "HIP SplitForward"); + + outputs.push_back(std::move(output)); + } + + return outputs; +} + +template +__global__ void SplitBackwardKernel(const T *const *grad_outputs, T *grad_input, int64_t N, int64_t H_in, int64_t W, + int64_t split_size, int64_t num_splits, const int64_t *H_outs) { + int64_t idx = blockIdx.x * blockDim.x + threadIdx.x; + int64_t total = N * H_in * W; + if (idx >= total) { + return; + } + + int64_t w = idx % W; + int64_t h = (idx / W) % H_in; + int64_t n = idx / (H_in * W); + + int64_t split_idx = h / split_size; + if (split_idx >= num_splits) { + return; + } + + int64_t H_out = H_outs[split_idx]; + int64_t local_h = h - split_idx * split_size; + + if (local_h >= H_out) { + return; + } + + const T *grad_output = grad_outputs[split_idx]; + T value = grad_output[(n * H_out + local_h) * W + w]; + grad_input[(n * H_in + h) * W + w] = value; +} + +template +std::shared_ptr LaunchSplitBackward(const std::vector &input_dims, int64_t split_size, int dim, + const std::vector> &grad_outputs) { + CHECK_GT(split_size, 0); + CHECK_GE(dim, 0) << "Currently we do not support negative dimension"; + CHECK_LT(dim, input_dims.size()); + + const auto &grad = grad_outputs[0]; + auto dtype = grad->Dtype(); + auto grad_input = std::make_shared(input_dims, dtype, grad->GetDevice()); + grad_input->Fill(0.0); + + int64_t N = std::accumulate(input_dims.begin(), input_dims.begin() + dim, 1, std::multiplies()); + int64_t W = std::accumulate(input_dims.begin() + dim + 1, input_dims.end(), 1, std::multiplies()); + int64_t H_in = input_dims[dim]; + int64_t num_splits = grad_outputs.size(); + + auto device = grad->GetDevice(); + const auto &stream = dynamic_cast( + infini_train::core::GetDeviceGuardImpl(device.type())->GetStream(device)) + ->hip_stream(); + // init the array of grad_output ptrs + std::vector host_grad_output_ptrs; + for (const auto &grad_output : grad_outputs) { + host_grad_output_ptrs.push_back(static_cast(grad_output->DataPtr())); + } + + void *device_ptr; + const T **device_grad_output_ptrs; + int64_t *device_H_outs; + hipMallocAsync(&device_ptr, (sizeof(T *) + sizeof(int64_t)) * num_splits, stream); + device_grad_output_ptrs = (const T **)(device_ptr); + device_H_outs = reinterpret_cast(device_grad_output_ptrs + num_splits); + + hipMemcpyAsync(device_grad_output_ptrs, host_grad_output_ptrs.data(), sizeof(T *) * num_splits, + hipMemcpyHostToDevice, stream); + + // init H_out for each split + std::vector H_outs(num_splits); + for (int i = 0; i < num_splits; ++i) { H_outs[i] = std::min(split_size, H_in - i * split_size); } + + hipMemcpyAsync(device_H_outs, H_outs.data(), sizeof(int64_t) * num_splits, hipMemcpyHostToDevice, stream); + + int64_t total_elements = N * H_in * W; + int threads_per_block = 256; + int num_blocks = (total_elements + threads_per_block - 1) / threads_per_block; + + SplitBackwardKernel<<>>(device_grad_output_ptrs, + static_cast(grad_input->DataPtr()), N, H_in, + W, split_size, num_splits, device_H_outs); + + HIP_CHECK(hipFree(device_ptr)); + + return grad_input; +} + +std::shared_ptr SplitBackward(const std::vector &input_dims, int64_t split_size, int dim, + const std::vector> &grad_outputs) { + CHECK_GT(split_size, 0); + CHECK_GE(dim, 0) << "Currently we do not support negative dimension"; + CHECK_LT(dim, input_dims.size()); + + return core::dcu::DispatchDcuFunc( + grad_outputs[0]->Dtype(), + [=]() { return LaunchSplitBackward(input_dims, split_size, dim, grad_outputs); }, + "HIP SplitBackward"); +} +} // namespace infini_train::kernels::dcu + +#define REGISTER_HIP_SPLIT_KERNEL(kernel_name) \ + REGISTER_KERNEL(infini_train::Device::DeviceType::kDCU, kernel_name, infini_train::kernels::dcu::kernel_name) + +REGISTER_HIP_SPLIT_KERNEL(SplitForward) +REGISTER_HIP_SPLIT_KERNEL(SplitBackward) + +#undef REGISTER_HIP_SPLIT_KERNEL diff --git a/infini_train/src/kernels/dcu/stack.hip b/infini_train/src/kernels/dcu/stack.hip new file mode 100644 index 00000000..378fd0ad --- /dev/null +++ b/infini_train/src/kernels/dcu/stack.hip @@ -0,0 +1,160 @@ +#include +#include +#include +#include +#include + +#include "glog/logging.h" + +#include "infini_train/include/common/dcu/common_dcu.h" +#include "infini_train/include/core/runtime/device_guard.h" +#include "infini_train/include/dispatcher.h" +#include "infini_train/include/tensor.h" + +#include "infini_train/src/core/runtime/dcu/dcu_dispatch.h" +#include "infini_train/src/core/runtime/dcu/dcu_runtime_common.h" + +namespace infini_train::kernels::dcu { +template +__global__ void StackForwardKernel(const T **inputs, T *output, int64_t N, int64_t D, int64_t num_inputs) { + int64_t idx = blockIdx.x * blockDim.x + threadIdx.x; + int64_t total = N * num_inputs * D; + + if (idx >= total) { + return; + } + + int64_t d = idx % D; + int64_t s = (idx / D) % num_inputs; + int64_t n = idx / (D * num_inputs); + + const T *input = inputs[s]; + output[idx] = input[n * D + d]; +} + +std::shared_ptr StackForward(const std::vector> &inputs, int64_t dim) { + CHECK(!inputs.empty()); + + const auto &base_dims = inputs[0]->Dims(); + auto dtype = inputs[0]->Dtype(); + if (dim < 0) { + dim += base_dims.size() + 1; + } + CHECK_GE(dim, 0); + CHECK_LE(dim, base_dims.size()); + for (const auto &input : inputs) { CHECK(input->Dims() == base_dims); } + + std::vector out_dims = base_dims; + out_dims.insert(out_dims.begin() + dim, inputs.size()); + auto output = std::make_shared(out_dims, dtype, inputs[0]->GetDevice()); + + const int64_t N = std::accumulate(base_dims.begin(), base_dims.begin() + dim, 1, std::multiplies()); + const int64_t D = std::accumulate(base_dims.begin() + dim, base_dims.end(), 1, std::multiplies()); + const int64_t num_inputs = inputs.size(); + + auto device = output->GetDevice(); + const auto &stream = dynamic_cast( + infini_train::core::GetDeviceGuardImpl(device.type())->GetStream(device)) + ->hip_stream(); + + int64_t total = N * num_inputs * D; + int threads_per_block = 256; + int num_blocks = (total + threads_per_block - 1) / threads_per_block; + + core::dcu::DispatchDcuFunc( + dtype, + [=]() { + std::vector host_input_ptrs; + for (const auto &t : inputs) { host_input_ptrs.push_back(static_cast(t->DataPtr())); } + + const T **device_input_ptrs; + hipMallocAsync(&device_input_ptrs, sizeof(T *) * num_inputs, stream); + hipMemcpyAsync(device_input_ptrs, host_input_ptrs.data(), sizeof(T *) * num_inputs, hipMemcpyHostToDevice, + stream); + + StackForwardKernel<<>>( + device_input_ptrs, static_cast(output->DataPtr()), N, D, num_inputs); + + HIP_CHECK(hipFree(device_input_ptrs)); + }, + "HIP StackForward"); + + return output; +} + +template +__global__ void StackBackwardKernel(const T *grad_output, T **grad_inputs, int64_t N, int64_t D, int64_t num_inputs) { + int64_t idx = blockIdx.x * blockDim.x + threadIdx.x; + int64_t total = N * num_inputs * D; + + if (idx >= total) { + return; + } + + int64_t d = idx % D; + int64_t s = (idx / D) % num_inputs; + int64_t n = idx / (D * num_inputs); + + if (s < num_inputs) { + grad_inputs[s][n * D + d] = grad_output[idx]; + } +} + +std::vector> StackBackward(const std::vector &input_dims, int64_t dim, + const std::shared_ptr &grad_output) { + if (dim < 0) { + dim += input_dims.size() + 1; + } + const int64_t num_inputs = grad_output->Dims()[dim]; + std::vector base_dims = grad_output->Dims(); + base_dims.erase(base_dims.begin() + dim); + + auto dtype = grad_output->Dtype(); + std::vector> grads; + for (int i = 0; i < num_inputs; ++i) { + auto t = std::make_shared(base_dims, dtype, grad_output->GetDevice()); + t->Fill(0.0); + grads.push_back(t); + } + + int64_t N = std::accumulate(input_dims.begin(), input_dims.begin() + dim, 1, std::multiplies()); + int64_t D = std::accumulate(input_dims.begin() + dim, input_dims.end(), 1, std::multiplies()); + + auto device = grad_output->GetDevice(); + const auto &stream = dynamic_cast( + infini_train::core::GetDeviceGuardImpl(device.type())->GetStream(device)) + ->hip_stream(); + + int64_t total = N * num_inputs * D; + int threads_per_block = 256; + int num_blocks = (total + threads_per_block - 1) / threads_per_block; + + core::dcu::DispatchDcuFunc( + dtype, + [=]() { + std::vector host_ptrs; + for (auto &t : grads) { host_ptrs.push_back(static_cast(t->DataPtr())); } + + T **device_ptrs; + hipMallocAsync(&device_ptrs, sizeof(T *) * num_inputs, stream); + hipMemcpyAsync(device_ptrs, host_ptrs.data(), sizeof(T *) * num_inputs, hipMemcpyHostToDevice, stream); + + StackBackwardKernel<<>>( + static_cast(grad_output->DataPtr()), device_ptrs, N, D, num_inputs); + + HIP_CHECK(hipFree(device_ptrs)); + }, + "HIP StackBackward"); + + return grads; +} + +} // namespace infini_train::kernels::dcu + +#define REGISTER_HIP_STACK_KERNEL(kernel_name) \ + REGISTER_KERNEL(infini_train::Device::DeviceType::kDCU, kernel_name, infini_train::kernels::dcu::kernel_name) + +REGISTER_HIP_STACK_KERNEL(StackForward) +REGISTER_HIP_STACK_KERNEL(StackBackward) + +#undef REGISTER_HIP_STACK_KERNEL diff --git a/infini_train/src/kernels/dcu/transform.hip b/infini_train/src/kernels/dcu/transform.hip new file mode 100644 index 00000000..1c963645 --- /dev/null +++ b/infini_train/src/kernels/dcu/transform.hip @@ -0,0 +1,593 @@ +#include +#include +#include +#include +#include + +#include "infini_train/include/common/dcu/common_dcu.h" +#include "infini_train/include/common/dcu/kernel_helper.cuh" +#include "infini_train/include/core/runtime/device_guard.h" +#include "infini_train/include/dispatcher.h" +#include "infini_train/include/tensor.h" + +#include "infini_train/src/core/runtime/dcu/dcu_dispatch.h" +#include "infini_train/src/core/runtime/dcu/dcu_runtime_common.h" + +namespace infini_train::kernels::dcu { + +template +__global__ void TrilForwardKernel(const T *input, T *output, int rows, int cols, int64_t diagonal) { + int idx = blockIdx.x * blockDim.x + threadIdx.x; + if (idx >= rows * cols) { + return; + } + + int row = idx / cols; + int col = idx % cols; + + if (row - col + diagonal >= 0) { + output[idx] = input[idx]; + } else { + output[idx] = T(0); + } +} + +std::shared_ptr TrilForward(const std::shared_ptr &input, int64_t diagonal) { + CHECK_EQ(input->Dims().size(), 2); + int64_t rows = input->Dims()[0]; + int64_t cols = input->Dims()[1]; + + auto output = std::make_shared(input->Dims(), input->Dtype(), input->GetDevice()); + + int threads_per_block = 256; + int num_blocks = (rows * cols + threads_per_block - 1) / threads_per_block; + + auto device = input->GetDevice(); + const auto &hip_stream = dynamic_cast( + infini_train::core::GetDeviceGuardImpl(device.type())->GetStream(device)) + ->hip_stream(); + + core::dcu::DispatchDcuFunc( + input->Dtype(), + [=]() { + TrilForwardKernel<<>>( + static_cast(input->DataPtr()), static_cast(output->DataPtr()), rows, cols, diagonal); + }, + "HIP TrilForward"); + + return output; +} + +template +__global__ void TrilBackwardKernel(const T *grad_output, T *grad_input, int rows, int cols, int64_t diagonal) { + int idx = blockIdx.x * blockDim.x + threadIdx.x; + if (idx >= rows * cols) { + return; + } + + int row = idx / cols; + int col = idx % cols; + + if (row - col + diagonal >= 0) { + grad_input[idx] = grad_output[idx]; + } else { + grad_input[idx] = T(0); + } +} + +std::shared_ptr TrilBackward(const std::shared_ptr &grad_output, int64_t diagonal) { + int rows = grad_output->Dims()[0]; + int cols = grad_output->Dims()[1]; + + auto dtype = grad_output->Dtype(); + auto grad_input = std::make_shared(grad_output->Dims(), dtype, grad_output->GetDevice()); + + int threads_per_block = 256; + int num_blocks = (rows * cols + threads_per_block - 1) / threads_per_block; + + auto device = grad_output->GetDevice(); + const auto &hip_stream = dynamic_cast( + infini_train::core::GetDeviceGuardImpl(device.type())->GetStream(device)) + ->hip_stream(); + + core::dcu::DispatchDcuFunc( + dtype, + [=]() { + grad_input->Fill(0.0); + TrilBackwardKernel<<>>( + static_cast(grad_output->DataPtr()), static_cast(grad_input->DataPtr()), rows, cols, + diagonal); + }, + "HIP TrilBackward"); + + return grad_input; +} + +template +__global__ void TriuForwardKernel(const T *input, T *output, int rows, int cols, int64_t diagonal) { + int idx = blockIdx.x * blockDim.x + threadIdx.x; + if (idx >= rows * cols) { + return; + } + + int row = idx / cols; + int col = idx % cols; + + if (row - col + diagonal <= 0) { + output[idx] = input[idx]; + } else { + output[idx] = T(0); + } +} + +std::shared_ptr TriuForward(const std::shared_ptr &input, int64_t diagonal) { + CHECK_EQ(input->Dims().size(), 2); + int64_t rows = input->Dims()[0]; + int64_t cols = input->Dims()[1]; + + auto output = std::make_shared(input->Dims(), input->Dtype(), input->GetDevice()); + + int threads_per_block = 256; + int num_blocks = (rows * cols + threads_per_block - 1) / threads_per_block; + + auto device = input->GetDevice(); + const auto &hip_stream = dynamic_cast( + infini_train::core::GetDeviceGuardImpl(device.type())->GetStream(device)) + ->hip_stream(); + + core::dcu::DispatchDcuFunc( + input->Dtype(), + [=]() { + TriuForwardKernel<<>>( + static_cast(input->DataPtr()), static_cast(output->DataPtr()), rows, cols, diagonal); + }, + "HIP TriuForward"); + + return output; +} + +template +__global__ void TriuBackwardKernel(const T *grad_output, T *grad_input, int rows, int cols, int64_t diagonal) { + int idx = blockIdx.x * blockDim.x + threadIdx.x; + if (idx >= rows * cols) { + return; + } + + int row = idx / cols; + int col = idx % cols; + + if (row - col + diagonal <= 0) { + grad_input[idx] = grad_output[idx]; + } else { + grad_input[idx] = T(0); + } +} + +std::shared_ptr TriuBackward(const std::shared_ptr &grad_output, int64_t diagonal) { + int rows = grad_output->Dims()[0]; + int cols = grad_output->Dims()[1]; + + auto dtype = grad_output->Dtype(); + auto grad_input = std::make_shared(grad_output->Dims(), dtype, grad_output->GetDevice()); + + int threads_per_block = 256; + int num_blocks = (rows * cols + threads_per_block - 1) / threads_per_block; + auto device = grad_output->GetDevice(); + const auto &hip_stream = dynamic_cast( + infini_train::core::GetDeviceGuardImpl(device.type())->GetStream(device)) + ->hip_stream(); + + core::dcu::DispatchDcuFunc( + dtype, + [=]() { + grad_input->Fill(0.0); + TriuBackwardKernel<<>>( + static_cast(grad_output->DataPtr()), static_cast(grad_input->DataPtr()), rows, cols, + diagonal); + }, + "HIP TriuBackward"); + + return grad_input; +} + +template +__global__ void TransposeForwardKernel(const T *input, T *output, const int64_t *in_dims, const int64_t *in_strides, + const int64_t *out_strides, int64_t ndim, int64_t dim0, int64_t dim1, + int64_t num_elements) { + int64_t idx = blockIdx.x * blockDim.x + threadIdx.x; + if (idx >= num_elements) { + return; + } + + int64_t remaining = idx; + // TODO(zbl): assume ndim <= 8 here + int64_t coords[8]; + + // 1. decode coord from output index + for (int i = 0; i < ndim; ++i) { + coords[i] = remaining / out_strides[i]; + remaining %= out_strides[i]; + } + + // 2. swap the coordinates + int64_t tmp = coords[dim0]; + coords[dim0] = coords[dim1]; + coords[dim1] = tmp; + + // 3. compute input flat index + int64_t in_flat_idx = 0; + for (int i = 0; i < ndim; ++i) { in_flat_idx += coords[i] * in_strides[i]; } + + output[idx] = input[in_flat_idx]; +} + +std::shared_ptr TransposeForward(const std::shared_ptr &input, int64_t dim0, int64_t dim1) { + // TODO(zbl): assume ndim <= 8 here + CHECK_LE(input->Dims().size(), 8); + dim0 = dim0 < 0 ? dim0 + input->Dims().size() : dim0; + dim1 = dim1 < 0 ? dim1 + input->Dims().size() : dim1; + CHECK(dim0 >= 0 && dim0 < input->Dims().size() && dim1 >= 0 && dim1 < input->Dims().size()); + + auto in_dims = input->Dims(); + std::vector out_dims = in_dims; + std::swap(out_dims[dim0], out_dims[dim1]); + + auto dtype = input->Dtype(); + auto output = std::make_shared(out_dims, dtype, input->GetDevice()); + int64_t ndim = in_dims.size(); + int64_t num_elements = output->NumElements(); + + // compute strides of in_dims and out_dims + std::vector in_strides(ndim, 1); + std::vector out_strides(ndim, 1); + for (int i = ndim - 2; i >= 0; --i) { + in_strides[i] = in_strides[i + 1] * in_dims[i + 1]; + out_strides[i] = out_strides[i + 1] * out_dims[i + 1]; + } + + auto device = input->GetDevice(); + const auto &stream = dynamic_cast( + infini_train::core::GetDeviceGuardImpl(device.type())->GetStream(device)) + ->hip_stream(); + + // Allocate device memory for dims and strides + // TODO(zbl): avoid using hipMalloc? + int64_t *device_buffer; + hipMallocAsync(&device_buffer, 3 * ndim * sizeof(int64_t), stream); + + int64_t *in_dims_dev = device_buffer; + int64_t *in_strides_dev = device_buffer + ndim; + int64_t *out_strides_dev = device_buffer + 2 * ndim; + + std::vector host_buffer; + host_buffer.insert(host_buffer.end(), in_dims.begin(), in_dims.end()); + host_buffer.insert(host_buffer.end(), in_strides.begin(), in_strides.end()); + host_buffer.insert(host_buffer.end(), out_strides.begin(), out_strides.end()); + + hipMemcpyAsync(device_buffer, host_buffer.data(), 3 * ndim * sizeof(int64_t), hipMemcpyHostToDevice, stream); + + int threads_per_block = 256; + int num_blocks = (num_elements + threads_per_block - 1) / threads_per_block; + + core::dcu::DispatchDcuFunc( + dtype, + [=]() { + output->Fill(0.0); + TransposeForwardKernel<<>>( + static_cast(input->DataPtr()), static_cast(output->DataPtr()), in_dims_dev, + in_strides_dev, out_strides_dev, ndim, dim0, dim1, num_elements); + }, + "HIP TransposeForward"); + + HIP_CHECK(hipFree(device_buffer)); + + return output; +} + +std::shared_ptr TransposeBackward(const std::shared_ptr &grad_output, int64_t dim0, int64_t dim1) { + return TransposeForward(grad_output, dim1, dim0); +} + +namespace { +enum class MaskMode { kLead, kTail }; + +static bool IsLeadMaskShape(const std::vector &in, const std::vector &mk) { + if (mk.empty() || in.empty()) { + return false; + } + if (mk.size() > in.size()) { + return false; + } + for (size_t d = 0; d < mk.size(); ++d) { + if (!(mk[d] == in[d] || mk[d] == 1)) { + return false; + } + } + return true; +} + +static bool IsTailMaskShape(const std::vector &in, const std::vector &mk) { + if (mk.size() > in.size()) { + return false; + } + size_t k = mk.size(); + for (size_t i = 0; i < k; ++i) { + int64_t in_dim = in[in.size() - k + i]; + int64_t mk_dim = mk[i]; + if (!(mk_dim == in_dim || mk_dim == 1)) { + return false; + } + } + return true; +} + +static MaskMode DecideMaskMode(const std::vector &in, const std::vector &mk) { + bool lead = IsLeadMaskShape(in, mk); + bool tail = IsTailMaskShape(in, mk); + CHECK(lead || tail) << "Mask must align/broadcast to either leading or trailing axes."; + // By default mask along tailing dims + return tail ? MaskMode::kTail : MaskMode::kLead; +} +} // namespace + +template +__global__ void MaskForwardKernel(const T *input, const T *mask, T *output, T value, int batch_size, int mask_size) { + int i = blockIdx.x * blockDim.x + threadIdx.x; + if (i < batch_size * mask_size) { + output[i] = (mask[i % mask_size] == T(1)) ? value : input[i]; + } +} + +template +__global__ void MaskLeadsForwardKernel(const T *input, const T *mask, T *output, T value, int rows, int inner) { + int i = blockIdx.x * blockDim.x + threadIdx.x; + if (i < rows * inner) { + output[i] = (mask[i / inner] == T(1)) ? value : input[i]; + } +} + +std::shared_ptr MaskForward(const std::shared_ptr &input, const std::shared_ptr &mask, + float value) { + auto input_shape = input->Dims(); + auto mask_shape = mask->Dims(); + auto dtype = input->Dtype(); + auto mask_casted = mask->Dtype() == dtype ? mask : std::make_shared(mask->To(dtype)); + // TODO(zbl): support bool mask + CHECK_EQ(static_cast(dtype), static_cast(mask_casted->Dtype())) + << "For now, input/mask dtypes must match."; + + MaskMode mode = DecideMaskMode(input_shape, mask_shape); + + auto output = std::make_shared(input_shape, dtype, input->GetDevice()); + auto device = output->GetDevice(); + const auto &hip_stream = dynamic_cast( + infini_train::core::GetDeviceGuardImpl(device.type())->GetStream(device)) + ->hip_stream(); + + int threads_per_block = 256; + + if (mode == MaskMode::kLead) { + int64_t rows = mask->NumElements(); + int64_t inner = input->NumElements() / rows; + int num_blocks = static_cast((input->NumElements() + threads_per_block - 1) / threads_per_block); + + core::dcu::DispatchDcuFunc( + dtype, + [=]() { + MaskLeadsForwardKernel<<>>( + static_cast(input->DataPtr()), static_cast(mask->DataPtr()), + static_cast(output->DataPtr()), common::dcu::Cast(value), rows, inner); + }, + "HIP MaskForward(rows)"); + } else { // kTail + int64_t mask_size = mask->NumElements(); + int64_t batch_size = input->NumElements() / mask_size; + int num_blocks = static_cast((input->NumElements() + threads_per_block - 1) / threads_per_block); + + core::dcu::DispatchDcuFunc( + dtype, + [=]() { + MaskForwardKernel<<>>( + static_cast(input->DataPtr()), static_cast(mask_casted->DataPtr()), + static_cast(output->DataPtr()), common::dcu::Cast(value), static_cast(batch_size), + static_cast(mask_size)); + }, + "HIP MaskForward(tail)"); + } + + return output; +} + +template +__global__ void MaskBackwardKernel(const T *grad_output, const T *mask, T *grad_input, int batch_size, int mask_size) { + int i = blockIdx.x * blockDim.x + threadIdx.x; + if (i < batch_size * mask_size) { + grad_input[i] = (mask[i % mask_size] == T(1)) ? T(0) : grad_output[i]; + } +} + +template +__global__ void MaskLeadsBackwardKernel(const T *grad_output, const T *mask, T *grad_input, int rows, int inner) { + int i = blockIdx.x * blockDim.x + threadIdx.x; + if (i < rows * inner) { + grad_input[i] = (mask[i / inner] == T(1)) ? T(0) : grad_output[i]; + } +} + +std::shared_ptr MaskBackward(const std::shared_ptr &grad_output, const std::shared_ptr &mask) { + auto output_shape = grad_output->Dims(); + auto mask_shape = mask->Dims(); + auto dtype = grad_output->Dtype(); + auto mask_casted = std::make_shared(mask->To(dtype)); + + MaskMode mode = DecideMaskMode(output_shape, mask_shape); + + auto grad_input = std::make_shared(output_shape, dtype, grad_output->GetDevice()); + auto device = grad_output->GetDevice(); + const auto &hip_stream = dynamic_cast( + infini_train::core::GetDeviceGuardImpl(device.type())->GetStream(device)) + ->hip_stream(); + + int threads_per_block = 256; + + if (mode == MaskMode::kLead) { + int64_t rows = mask->NumElements(); + int64_t inner = grad_output->NumElements() / rows; + int num_blocks = static_cast((grad_output->NumElements() + threads_per_block - 1) / threads_per_block); + + core::dcu::DispatchDcuFunc( + dtype, + [=]() { + grad_input->Fill(0.0); + MaskLeadsBackwardKernel<<>>( + static_cast(grad_output->DataPtr()), static_cast(mask_casted->DataPtr()), + static_cast(grad_input->DataPtr()), rows, inner); + }, + "HIP MaskBackward(rows)"); + } else { // kTail + int64_t mask_size = mask->NumElements(); + int64_t batch_size = grad_output->NumElements() / mask_size; + int num_blocks = static_cast((grad_output->NumElements() + threads_per_block - 1) / threads_per_block); + + core::dcu::DispatchDcuFunc( + dtype, + [=]() { + grad_input->Fill(0.0); + MaskBackwardKernel<<>>( + static_cast(grad_output->DataPtr()), static_cast(mask_casted->DataPtr()), + static_cast(grad_input->DataPtr()), static_cast(batch_size), static_cast(mask_size)); + }, + "HIP MaskBackward(tail)"); + } + + return grad_input; +} + +template +__global__ void RepeatInterleaveForwardKernel(const T *input, T *output, int64_t outer, int64_t dim_size, int64_t inner, + int64_t repeat) { + int64_t idx = blockIdx.x * blockDim.x + threadIdx.x; + int64_t total = outer * dim_size * repeat * inner; + if (idx >= total) { + return; + } + + int64_t i = idx / inner; + int64_t j = idx % inner; + + int64_t o = i / (dim_size * repeat); + int64_t di = (i / repeat) % dim_size; + + output[idx] = input[(o * dim_size + di) * inner + j]; +} + +std::shared_ptr RepeatInterleaveForward(const std::shared_ptr &input, int64_t repeat, int64_t dim) { + CHECK_GT(repeat, 0); + CHECK_GE(dim, 0); + CHECK_LT(dim, input->Dims().size()); + + const auto &input_dims = input->Dims(); + const int64_t outer = std::accumulate(input_dims.begin(), input_dims.begin() + dim, 1, std::multiplies()); + const int64_t inner + = std::accumulate(input_dims.begin() + dim + 1, input_dims.end(), 1, std::multiplies()); + const int64_t dim_size = input_dims[dim]; + + std::vector output_dims = input_dims; + output_dims[dim] = dim_size * repeat; + auto output = std::make_shared(output_dims, input->Dtype(), input->GetDevice()); + + int64_t total_elements = outer * dim_size * repeat * inner; + int threads_per_block = 256; + int num_blocks = (total_elements + threads_per_block - 1) / threads_per_block; + auto device = input->GetDevice(); + const auto &hip_stream = dynamic_cast( + infini_train::core::GetDeviceGuardImpl(device.type())->GetStream(device)) + ->hip_stream(); + + core::dcu::DispatchDcuFunc( + input->Dtype(), + [=]() { + RepeatInterleaveForwardKernel<<>>( + static_cast(input->DataPtr()), static_cast(output->DataPtr()), outer, dim_size, inner, + repeat); + }, + "HIP RepeatInterleaveForward"); + + return output; +} + +template +__global__ void RepeatInterleaveBackwardKernel(const T *grad_output, T *grad_input, int64_t outer, int64_t dim_size, + int64_t inner, int64_t repeat) { + int64_t idx = blockIdx.x * blockDim.x + threadIdx.x; + int64_t total = outer * dim_size * inner; + if (idx >= total) { + return; + } + + int64_t i = idx / inner; + int64_t j = idx % inner; + + int64_t o = i / dim_size; + int64_t di = i % dim_size; + + T sum = T(0); + for (int64_t r = 0; r < repeat; ++r) { + int64_t out_idx = ((o * dim_size * repeat + di * repeat + r) * inner) + j; + sum += grad_output[out_idx]; + } + grad_input[idx] = sum; +} + +std::shared_ptr RepeatInterleaveBackward(const std::shared_ptr &grad_output, + const std::vector &input_dims, int64_t dim) { + CHECK_GE(dim, 0); + CHECK_LT(dim, input_dims.size()); + + const int64_t outer = std::accumulate(input_dims.begin(), input_dims.begin() + dim, 1, std::multiplies()); + const int64_t inner + = std::accumulate(input_dims.begin() + dim + 1, input_dims.end(), 1, std::multiplies()); + const int64_t dim_size = input_dims[dim]; + + int64_t repeat = grad_output->Dims()[dim] / dim_size; + CHECK_EQ(grad_output->Dims()[dim], dim_size * repeat); + + auto grad_input = std::make_shared(input_dims, grad_output->Dtype(), grad_output->GetDevice()); + + int64_t total_elements = outer * dim_size * inner; + int threads_per_block = 256; + int num_blocks = (total_elements + threads_per_block - 1) / threads_per_block; + auto device = grad_output->GetDevice(); + const auto &hip_stream = dynamic_cast( + infini_train::core::GetDeviceGuardImpl(device.type())->GetStream(device)) + ->hip_stream(); + + core::dcu::DispatchDcuFunc( + grad_output->Dtype(), + [=]() { + grad_input->Fill(0.0); + RepeatInterleaveBackwardKernel<<>>( + static_cast(grad_output->DataPtr()), static_cast(grad_input->DataPtr()), outer, + dim_size, inner, repeat); + }, + "HIP RepeatInterleaveBackward"); + + return grad_input; +} +} // namespace infini_train::kernels::dcu + +#define REGISTER_HIP_TRANSFORM_KERNEL(kernel_name) \ + REGISTER_KERNEL(infini_train::Device::DeviceType::kDCU, kernel_name, infini_train::kernels::dcu::kernel_name) + +REGISTER_HIP_TRANSFORM_KERNEL(TrilForward) +REGISTER_HIP_TRANSFORM_KERNEL(TrilBackward) +REGISTER_HIP_TRANSFORM_KERNEL(TriuForward) +REGISTER_HIP_TRANSFORM_KERNEL(TriuBackward) +REGISTER_HIP_TRANSFORM_KERNEL(TransposeForward) +REGISTER_HIP_TRANSFORM_KERNEL(TransposeBackward) +REGISTER_HIP_TRANSFORM_KERNEL(MaskForward) +REGISTER_HIP_TRANSFORM_KERNEL(MaskBackward) +REGISTER_HIP_TRANSFORM_KERNEL(RepeatInterleaveForward) +REGISTER_HIP_TRANSFORM_KERNEL(RepeatInterleaveBackward) + +#undef REGISTER_HIP_TRANSFORM_KERNEL diff --git a/infini_train/src/kernels/dcu/vocab_parallel_cross_entropy.hip b/infini_train/src/kernels/dcu/vocab_parallel_cross_entropy.hip new file mode 100644 index 00000000..85a0792b --- /dev/null +++ b/infini_train/src/kernels/dcu/vocab_parallel_cross_entropy.hip @@ -0,0 +1,127 @@ +#include + +#include + +#include "infini_train/include/common/dcu/common_dcu.h" +#include "infini_train/include/common/dcu/kernel_helper.cuh" +#include "infini_train/include/core/runtime/device_guard.h" +#include "infini_train/include/dispatcher.h" +#include "infini_train/include/tensor.h" + +#include "infini_train/src/core/runtime/dcu/dcu_dispatch.h" +#include "infini_train/src/core/runtime/dcu/dcu_runtime_common.h" + +namespace infini_train::kernels::dcu { + +template +__global__ void +VocabParallelCrossEntropyBackwardKernel(const Tinput *__restrict__ softmax_local, // [rows, V_local] + Tinput *__restrict__ grad_input, // [rows, V_local] + const Tindex *__restrict__ masked_target, // [rows] + const Tmask *__restrict__ target_mask_row, // [rows]，0/1 + const Tmask *__restrict__ valid_mask_local, // [rows, V_local] or [1, V_local] + const Tinput *__restrict__ dloss_buf, // [1] or [rows] + int rows, int V_local, + int dloss_is_scalar, // 1=scalaer，0=by row + float one_minus_label_smoothing, // 1 - label_smoothing + float smoothing_term // label_smoothing / vocab_size_original +) { + const int r = blockIdx.x; + if (r >= rows) { + return; + } + + const float dm = common::dcu::Cast(dloss_is_scalar ? dloss_buf[0] : dloss_buf[r]); + const float vm_row = 1.0f - common::dcu::Cast(target_mask_row[r]); + const float row_scale = dm * one_minus_label_smoothing * vm_row; + const Tindex t = masked_target[r]; + + for (int j = threadIdx.x; j < V_local; j += BLOCK_SIZE) { + const int idx = r * V_local + j; + + const float s = common::dcu::Cast(softmax_local[idx]); + const float vm = common::dcu::Cast(valid_mask_local[j]); + + float grad = dm * s; + + if (static_cast(t) >= 0 && j == static_cast(t)) { + grad -= row_scale; + } + + grad -= dm * smoothing_term * vm; + grad *= vm; + + grad_input[idx] = common::dcu::Cast(grad); + } +} + +std::shared_ptr +VocabParallelCrossEntropyBackward(const std::shared_ptr &grad_output, // [rows] + const std::shared_ptr &softmax_local, // [rows, V_local] + const std::shared_ptr &target_mask, // [rows] + const std::shared_ptr &masked_target, // [rows]，int64 + const std::shared_ptr &valid_mask_local, // [1, V_local] + const int64_t vocab_size_local, const int64_t vocab_size_original, + float label_smoothing) { + + const int64_t rows = softmax_local->NumElements() / vocab_size_local; + CHECK_EQ(masked_target->NumElements(), rows); + CHECK_EQ(target_mask->NumElements(), rows); + CHECK_EQ(valid_mask_local->NumElements(), vocab_size_local); + + int dloss_is_scalar = 0; + if (grad_output->Dims().size() == 0) { + dloss_is_scalar = 1; + } else { + CHECK(grad_output->NumElements() == rows || grad_output->NumElements() == 1) + << "grad_output must be scalar or length rows"; + dloss_is_scalar = (grad_output->NumElements() == 1); + } + + auto device = grad_output->GetDevice(); + const auto &hip_stream = dynamic_cast( + infini_train::core::GetDeviceGuardImpl(device.type())->GetStream(device)) + ->hip_stream(); + + // logits should be [rows, V_local] + auto grad_input = std::make_shared(softmax_local->Dims(), softmax_local->Dtype(), device); + + const float one_minus_label_smoothing = 1.0f - label_smoothing; + const float smoothing_term = (label_smoothing > 0.f && vocab_size_original > 0) + ? (label_smoothing / static_cast(vocab_size_original)) + : 0.0f; + + constexpr int threads_per_block = 256; + const int num_blocks = static_cast(rows); + + core::dcu::DispatchDcuFunc, + DataTypeList>( + {masked_target->Dtype(), softmax_local->Dtype()}, + [=]() { + using Tmask = Tinput; + + const Tinput *softmax_ptr = static_cast(softmax_local->DataPtr()); + const Tmask *tmask_ptr = static_cast(target_mask->DataPtr()); + const Tmask *vml_ptr = static_cast(valid_mask_local->DataPtr()); + const Tindex *mtarget_ptr = static_cast(masked_target->DataPtr()); + const Tinput *grad_output_ptr = static_cast(grad_output->DataPtr()); + Tinput *grad_input_ptr = static_cast(grad_input->DataPtr()); + + VocabParallelCrossEntropyBackwardKernel + <<>>(softmax_ptr, grad_input_ptr, mtarget_ptr, tmask_ptr, + vml_ptr, grad_output_ptr, static_cast(rows), + static_cast(vocab_size_local), dloss_is_scalar, + one_minus_label_smoothing, smoothing_term); + }, + "HIP VocabParallelCrossEntropyBackward"); + + return grad_input; +} +} // namespace infini_train::kernels::dcu + +#define REGISTER_HIP_VOCAB_PARALLEL_CROSS_ENTROPY_KERNEL(kernel_name) \ + REGISTER_KERNEL(infini_train::Device::DeviceType::kDCU, kernel_name, infini_train::kernels::dcu::kernel_name) + +REGISTER_HIP_VOCAB_PARALLEL_CROSS_ENTROPY_KERNEL(VocabParallelCrossEntropyBackward) + +#undef REGISTER_HIP_CROSS_ENTROPY_KERNEL diff --git a/infini_train/src/nn/parallel/data_parallel.cc b/infini_train/src/nn/parallel/data_parallel.cc index c48761b7..29953c9f 100644 --- a/infini_train/src/nn/parallel/data_parallel.cc +++ b/infini_train/src/nn/parallel/data_parallel.cc @@ -1,6 +1,5 @@ #include "infini_train/include/nn/parallel/data_parallel.h" -#include #include #include #include @@ -22,8 +21,8 @@ constexpr char kModuleName[] = "module"; std::vector>> ParallelApply(const std::vector> &modules, const std::vector>> &inputs, const std::vector &devices) { - CHECK_EQ(modules.size(), inputs.size()) << std::format( - "The number of modules {} is not equal to the number of inputs {}", modules.size(), inputs.size()); + CHECK_EQ(modules.size(), inputs.size()) << "The number of modules " << modules.size() + << " is not equal to the number of inputs " << inputs.size(); CHECK_EQ(modules.size(), devices.size()); // pre-allocate results so we do not need lock in the worker threads @@ -80,9 +79,8 @@ std::vector> DataParallel::Forward(const std::vectorParameters()) { if (tensor->GetDevice() != src_device_) { - LOG(FATAL) << std::format("module must have its Parameters on device {} (device_ids[0]) but found " - "one of them on device: {}", - src_device_.ToString(), tensor->GetDevice().ToString()); + LOG(FATAL) << "module must have its Parameters on device " << src_device_.ToString() + << " (device_ids[0]) but found one of them on device: " << tensor->GetDevice().ToString(); } } diff --git a/infini_train/src/nn/parallel/global.cc b/infini_train/src/nn/parallel/global.cc index 65a3208e..2b868b20 100644 --- a/infini_train/src/nn/parallel/global.cc +++ b/infini_train/src/nn/parallel/global.cc @@ -1,7 +1,7 @@ #include "infini_train/include/nn/parallel/global.h" #include -#include +#include #include #include "glog/logging.h" @@ -195,9 +195,9 @@ inline int NumGroups(const Layout &L, Axis target) { std::string ProcessGroupOverview(const Layout &L, bool skip_trivial_axes) { std::ostringstream oss; - oss << std::format("\n=== Parallel Communication Groups ===\n" - "world_size = {}, config: {{DP={}, TP={}, PP={}}}, order: {{", - GetWorldSize(), L.sizes[DP], L.sizes[TP], L.sizes[PP]); + oss << "\n=== Parallel Communication Groups ===\n" + << "world_size = " << GetWorldSize() << ", config: {DP=" << L.sizes[DP] << ", TP=" << L.sizes[TP] + << ", PP=" << L.sizes[PP] << "}, order: {"; for (int i = 0; i < AXIS_COUNT; ++i) { oss << AxisName(L.order[i]) << (i + 1 == AXIS_COUNT ? "" : " -> "); } oss << "}\n"; @@ -205,7 +205,7 @@ std::string ProcessGroupOverview(const Layout &L, bool skip_trivial_axes) { for (int a = 0; a < AXIS_COUNT; ++a) { Axis ax = static_cast(a); if (skip_trivial_axes && L.sizes[ax] <= 1) { - oss << std::format("[{}] size={}, unenabled\n", AxisName(ax), L.sizes[ax]); + oss << "[" << AxisName(ax) << "] size=" << L.sizes[ax] << ", unenabled\n"; continue; } // Build > mapping @@ -223,7 +223,7 @@ std::string ProcessGroupOverview(const Layout &L, bool skip_trivial_axes) { const int num_groups = NumGroups(L, ax); const auto name = AxisName(ax); - oss << std::format("[{}] size={}, num_groups={}\n", name, L.sizes[ax], num_groups); + oss << "[" << name << "] size=" << L.sizes[ax] << ", num_groups=" << num_groups << "\n"; // Iterate and print in the order of Group ID for (const auto &pair : groups) { @@ -245,8 +245,8 @@ std::string ProcessGroupOverview(const Layout &L, bool skip_trivial_axes) { } ranks_str += std::to_string(ranks[i]); } - oss << std::format(" - {} {} (dp={}, tp={}, pp={}): [{}]\n", name, gid, dp_size_str, tp_size_str, - pp_size_str, ranks_str); + oss << " - " << name << " " << gid << " (dp=" << dp_size_str << ", tp=" << tp_size_str + << ", pp=" << pp_size_str << "): [" << ranks_str << "]\n"; } if (a + 1 < AXIS_COUNT) { oss << "\n"; diff --git a/infini_train/src/nn/parallel/pp/pipeline_schedule.cc b/infini_train/src/nn/parallel/pp/pipeline_schedule.cc index 090b7b15..0826b289 100644 --- a/infini_train/src/nn/parallel/pp/pipeline_schedule.cc +++ b/infini_train/src/nn/parallel/pp/pipeline_schedule.cc @@ -2,7 +2,9 @@ #include "infini_train/include/nn/parallel/pp/pipeline_schedule.h" #include +#include #include +#include #include #include "glog/logging.h" @@ -24,9 +26,9 @@ void PrintScheduleTable(const std::vector &sche int vpp_size) { int total_global_chunks = num_stages * vpp_size; - LOG(INFO) << std::format("=== Schedule Table ===\n" - "n: {}, stages: {}, vpp: {}, total_chunks: {}", - n, num_stages, vpp_size, total_global_chunks); + LOG(INFO) << "=== Schedule Table ===\n" + << "n: " << n << ", stages: " << num_stages << ", vpp: " << vpp_size + << ", total_chunks: " << total_global_chunks; LOG(INFO) << ""; LOG(INFO) << "Step | Type | Microbatch | Global Chunk | Local Chunk | Stage"; LOG(INFO) << "-----|-----------|------------|--------------|-------------|-------"; @@ -37,9 +39,11 @@ void PrintScheduleTable(const std::vector &sche std::string type_str = task.is_forward ? "Forward" : "Backward"; - auto s_info = std::format("{:4} | {:<9} | {:>10} | {:>12} | {:>11} | {:>5}", task.step, type_str, - task.microbatch_id, task.global_chunk_id, local_chunk, owning_stage); - LOG(INFO) << s_info; + std::ostringstream s_info; + s_info << std::setw(4) << task.step << " | " << std::left << std::setw(9) << type_str << std::right << " | " + << std::setw(10) << task.microbatch_id << " | " << std::setw(12) << task.global_chunk_id << " | " + << std::setw(11) << local_chunk << " | " << std::setw(5) << owning_stage; + LOG(INFO) << s_info.str(); } } diff --git a/infini_train/src/profiler.cc b/infini_train/src/profiler.cc index d53cc351..8265fa57 100644 --- a/infini_train/src/profiler.cc +++ b/infini_train/src/profiler.cc @@ -2,8 +2,10 @@ #include #include +#include #include #include +#include #include "glog/logging.h" @@ -235,7 +237,9 @@ void Profiler::Report(const std::string &file_prefix, SortBy sort_by) const { auto get_stream = [&](int64_t rank) -> std::ostream & { auto &file = file_map[rank]; if (!file.is_open()) { - std::string filename = std::format("{}.rank{}", file_prefix, rank); + std::ostringstream filename_stream; + filename_stream << file_prefix << ".rank" << rank; + std::string filename = filename_stream.str(); file.open(filename); if (!file) { LOG(ERROR) << "Failed to open file: " << filename; @@ -298,7 +302,9 @@ void Profiler::PrintRecords(const std::string &file_prefix) const { auto get_stream = [&](int64_t rank) -> std::ostream & { auto &file = file_map[rank]; if (!file.is_open()) { - std::string filename = std::format("{}.rank{}", file_prefix, rank); + std::ostringstream filename_stream; + filename_stream << file_prefix << ".rank" << rank; + std::string filename = filename_stream.str(); file.open(filename); if (!file) { LOG(ERROR) << "Failed to open file: " << filename; diff --git a/infini_train/src/tensor.cc b/infini_train/src/tensor.cc index 18ca3d22..a54edc1a 100644 --- a/infini_train/src/tensor.cc +++ b/infini_train/src/tensor.cc @@ -3,6 +3,7 @@ #include #include #include +#include #include #include From 3ac18a6b74c1c87c18ed669612d288c94bbcaba8 Mon Sep 17 00:00:00 2001 From: flyingdown Date: Wed, 5 Aug 2026 12:35:11 +0800 Subject: [PATCH 2/4] style: format DCU backend changes --- infini_train/include/common/dcu/common_dcu.h | 18 +++++++++--------- .../include/common/dcu/kernel_helper.cuh | 3 +-- infini_train/include/common/dcu/rccl_compat.h | 13 +++---------- .../src/core/runtime/dcu/dcu_dispatch.h | 2 +- .../src/core/runtime/dcu/dcu_runtime_common.cc | 4 +--- infini_train/src/nn/parallel/data_parallel.cc | 4 ++-- 6 files changed, 17 insertions(+), 27 deletions(-) diff --git a/infini_train/include/common/dcu/common_dcu.h b/infini_train/include/common/dcu/common_dcu.h index 6c200f2f..7a8aeb56 100644 --- a/infini_train/include/common/dcu/common_dcu.h +++ b/infini_train/include/common/dcu/common_dcu.h @@ -15,20 +15,20 @@ namespace infini_train::common::dcu { // Common HIP Macros -#define HIP_CHECK(call) \ +#define HIP_CHECK(call) \ do { \ - hipError_t status = call; \ - if (status != hipSuccess) { \ - LOG(FATAL) << "HIP Error: " << hipGetErrorString(status) << " at " << __FILE__ << ":" << __LINE__; \ + hipError_t status = call; \ + if (status != hipSuccess) { \ + LOG(FATAL) << "HIP Error: " << hipGetErrorString(status) << " at " << __FILE__ << ":" << __LINE__; \ } \ } while (0) -#define HIPBLAS_CHECK(call) \ +#define HIPBLAS_CHECK(call) \ do { \ - hipblasStatus_t status = call; \ - if (status != HIPBLAS_STATUS_SUCCESS) { \ - LOG(FATAL) << "HIPBLAS Error: status=" << static_cast(status) << " at " << __FILE__ << ":" \ - << __LINE__; \ + hipblasStatus_t status = call; \ + if (status != HIPBLAS_STATUS_SUCCESS) { \ + LOG(FATAL) << "HIPBLAS Error: status=" << static_cast(status) << " at " << __FILE__ << ":" \ + << __LINE__; \ } \ } while (0) diff --git a/infini_train/include/common/dcu/kernel_helper.cuh b/infini_train/include/common/dcu/kernel_helper.cuh index 262fd253..29b13e6d 100644 --- a/infini_train/include/common/dcu/kernel_helper.cuh +++ b/infini_train/include/common/dcu/kernel_helper.cuh @@ -28,8 +28,7 @@ __device__ __forceinline__ void AtomicAdd(hip_bfloat16 *address, hip_bfloat16 va hip_bfloat16 updated(static_cast(current) + static_cast(value)); const unsigned int updated_bits = static_cast(updated.data); const unsigned int replacement - = upper ? ((assumed & 0x0000ffffU) | (updated_bits << 16)) - : ((assumed & 0xffff0000U) | updated_bits); + = upper ? ((assumed & 0x0000ffffU) | (updated_bits << 16)) : ((assumed & 0xffff0000U) | updated_bits); old = atomicCAS(base, assumed, replacement); } while (old != assumed); } diff --git a/infini_train/include/common/dcu/rccl_compat.h b/infini_train/include/common/dcu/rccl_compat.h index b1e418f0..d2d3f5d7 100644 --- a/infini_train/include/common/dcu/rccl_compat.h +++ b/infini_train/include/common/dcu/rccl_compat.h @@ -44,14 +44,7 @@ typedef enum { ncclNumTypes = 10 } ncclDataType_t; -typedef enum { - ncclSum = 0, - ncclProd = 1, - ncclMax = 2, - ncclMin = 3, - ncclAvg = 4, - ncclNumOps = 5 -} ncclRedOp_t; +typedef enum { ncclSum = 0, ncclProd = 1, ncclMax = 2, ncclMin = 3, ncclAvg = 4, ncclNumOps = 5 } ncclRedOp_t; extern "C" { const char *ncclGetErrorString(ncclResult_t result); @@ -62,8 +55,8 @@ ncclResult_t ncclCommDestroy(ncclComm_t comm); ncclResult_t ncclCommGetAsyncError(ncclComm_t comm, ncclResult_t *asyncError); ncclResult_t ncclGroupStart(); ncclResult_t ncclGroupEnd(); -ncclResult_t ncclAllReduce(const void *sendbuff, void *recvbuff, size_t count, ncclDataType_t datatype, - ncclRedOp_t op, ncclComm_t comm, hipStream_t stream); +ncclResult_t ncclAllReduce(const void *sendbuff, void *recvbuff, size_t count, ncclDataType_t datatype, ncclRedOp_t op, + ncclComm_t comm, hipStream_t stream); ncclResult_t ncclBroadcast(const void *sendbuff, void *recvbuff, size_t count, ncclDataType_t datatype, int root, ncclComm_t comm, hipStream_t stream); ncclResult_t ncclReduce(const void *sendbuff, void *recvbuff, size_t count, ncclDataType_t datatype, ncclRedOp_t op, diff --git a/infini_train/src/core/runtime/dcu/dcu_dispatch.h b/infini_train/src/core/runtime/dcu/dcu_dispatch.h index 8bc2767f..1159817c 100644 --- a/infini_train/src/core/runtime/dcu/dcu_dispatch.h +++ b/infini_train/src/core/runtime/dcu/dcu_dispatch.h @@ -43,7 +43,7 @@ auto DispatchDcuFunc(DataType dtype, Functor &&func, std::string_view context_id template auto DispatchDcuFunc(const std::vector &dtypes, Functor &&func, std::string_view context_identifier = "", - Args &&...args) { + Args &&...args) { return infini_train::DispatchByTypeMap( dtypes, std::forward(func), context_identifier, std::forward(args)...); } diff --git a/infini_train/src/core/runtime/dcu/dcu_runtime_common.cc b/infini_train/src/core/runtime/dcu/dcu_runtime_common.cc index 273bf9fd..641eef0a 100644 --- a/infini_train/src/core/runtime/dcu/dcu_runtime_common.cc +++ b/infini_train/src/core/runtime/dcu/dcu_runtime_common.cc @@ -34,9 +34,7 @@ hipEvent_t DcuEvent::hip_event() const { return event_; } DcuStream::DcuStream() { HIP_CHECK(hipStreamCreate(&stream_)); } -DcuStream::DcuStream(int priority) { - HIP_CHECK(hipStreamCreateWithPriority(&stream_, hipStreamNonBlocking, priority)); -} +DcuStream::DcuStream(int priority) { HIP_CHECK(hipStreamCreateWithPriority(&stream_, hipStreamNonBlocking, priority)); } DcuStream::~DcuStream() { // Do nothing. diff --git a/infini_train/src/nn/parallel/data_parallel.cc b/infini_train/src/nn/parallel/data_parallel.cc index 29953c9f..3adc1aa6 100644 --- a/infini_train/src/nn/parallel/data_parallel.cc +++ b/infini_train/src/nn/parallel/data_parallel.cc @@ -21,8 +21,8 @@ constexpr char kModuleName[] = "module"; std::vector>> ParallelApply(const std::vector> &modules, const std::vector>> &inputs, const std::vector &devices) { - CHECK_EQ(modules.size(), inputs.size()) << "The number of modules " << modules.size() - << " is not equal to the number of inputs " << inputs.size(); + CHECK_EQ(modules.size(), inputs.size()) + << "The number of modules " << modules.size() << " is not equal to the number of inputs " << inputs.size(); CHECK_EQ(modules.size(), devices.size()); // pre-allocate results so we do not need lock in the worker threads From e004ad65e6efb7580a2ca81220d43db1764f5dc5 Mon Sep 17 00:00:00 2001 From: cy <2833839179@qq.com> Date: Wed, 5 Aug 2026 13:50:53 +0800 Subject: [PATCH 3/4] fix: resolve DCU-only PR conflicts --- CMakeLists.txt | 219 ++++++++++------------------------ README.md | 13 +- example/gpt2/main.cc | 47 ++------ example/llama3/main.cc | 47 ++------ infini_train/include/device.h | 13 +- infini_train/src/device.cc | 11 -- 6 files changed, 91 insertions(+), 259 deletions(-) diff --git a/CMakeLists.txt b/CMakeLists.txt index 21098ad6..72d0809f 100644 --- a/CMakeLists.txt +++ b/CMakeLists.txt @@ -1,25 +1,12 @@ -<<<<<<< ours cmake_minimum_required(VERSION 3.28) -======= -# Platforms -option(USE_CUDA "Support NVIDIA CUDA" OFF) -option(USE_MACA "Support MetaX MACA" OFF) -option(USE_DCU "Support Hygon DCU through DTK/HIP" OFF) ->>>>>>> theirs option(USE_CUDA "Support NVIDIA CUDA" OFF) +option(USE_DCU "Support Hygon DCU through DTK/HIP" OFF) option(PROFILE_MODE "ENABLE PROFILE MODE" OFF) option(USE_OMP "Use OpenMP as backend for Eigen" ON) -<<<<<<< ours option(USE_NCCL "Build project for distributed running" ON) -option(BUILD_TEST "Build InfiniTrain tests" OFF) -======= -option(USE_NCCL "Build project for distributed running on CUDA using NCCL" ON) -option(USE_MCCL "Build project for distributed running on MACA using MCCL" ON) option(USE_RCCL "Build project for distributed running on DCU using RCCL" OFF) -option(USE_MPI "Enable MPI for inter-node CPU communication" ON) -cmake_minimum_required(VERSION 3.28) ->>>>>>> theirs +option(BUILD_TEST "Build InfiniTrain tests" OFF) project(infini_train VERSION 0.6.0 LANGUAGES CXX) @@ -87,9 +74,9 @@ if(NOT USE_NCCL) list(FILTER SRC EXCLUDE REGEX ".*infini_train/src/core/ccl/cuda/.*") endif() if(NOT USE_DCU) - list(FILTER SRC EXCLUDE REGEX ".*/(ccl|runtime)/dcu/.*") + list(FILTER SRC EXCLUDE REGEX ".*/(ccl|runtime)/dcu/.*") elseif(NOT USE_RCCL) - list(FILTER SRC EXCLUDE REGEX ".*/ccl/dcu/.*") + list(FILTER SRC EXCLUDE REGEX ".*/ccl/dcu/.*") endif() # CPU kernels (*.cc) @@ -143,11 +130,54 @@ if(USE_CUDA) endif() endif() +# ------------------------------------------------------------------------------ +# DCU kernels library (optional) +# ------------------------------------------------------------------------------ + +if(USE_DCU) + add_compile_definitions(USE_DCU=1) + + set(DCU_PATH "$ENV{DTK_PATH}" CACHE PATH "Hygon DTK installation root") + set(DCU_ARCH "" CACHE STRING "Optional HIP offload architecture reported by rocminfo") + if(NOT DCU_PATH) + set(DCU_PATH /opt/dtk) + endif() + + find_program(HIPCC_EXECUTABLE hipcc + HINTS "${DCU_PATH}/bin" "${DCU_PATH}/llvm/bin" /opt/rocm/bin + REQUIRED) + if(NOT CMAKE_CXX_COMPILER MATCHES "hipcc") + message(WARNING + "DCU kernels must be compiled by hipcc. Reconfigure with " + "-DCMAKE_CXX_COMPILER=${HIPCC_EXECUTABLE}") + endif() + + include_directories("${DCU_PATH}/include") + link_directories("${DCU_PATH}/lib" "${DCU_PATH}/lib64") + + find_library(DCU_RUNTIME_LIB NAMES amdhip64 hip_hcc + HINTS "${DCU_PATH}/lib" "${DCU_PATH}/lib64" /opt/rocm/lib /opt/rocm/lib64 + REQUIRED) + find_library(DCU_BLAS_LIB NAMES hipblas + HINTS "${DCU_PATH}/lib" "${DCU_PATH}/lib64" /opt/rocm/lib /opt/rocm/lib64 + REQUIRED) + + file(GLOB_RECURSE DCU_KERNELS ${PROJECT_SOURCE_DIR}/infini_train/src/kernels/dcu/*.hip) + set_source_files_properties(${DCU_KERNELS} PROPERTIES + LANGUAGE CXX + COMPILE_OPTIONS "-x;hip" + ) + add_library(infini_train_dcu_kernels STATIC ${DCU_KERNELS}) + if(DCU_ARCH) + target_compile_options(infini_train_dcu_kernels PRIVATE "--offload-arch=${DCU_ARCH}") + endif() + target_link_libraries(infini_train_dcu_kernels PUBLIC glog ${DCU_RUNTIME_LIB} ${DCU_BLAS_LIB}) +endif() + # ------------------------------------------------------------------------------ # Main framework library # ------------------------------------------------------------------------------ -<<<<<<< ours add_library(infini_train STATIC ${SRC}) target_link_libraries(infini_train PUBLIC @@ -172,128 +202,24 @@ if(USE_CUDA) # keep this. Otherwise it's harmless. target_link_libraries(infini_train PUBLIC nccl) endif() -======= -# ========================= -# DCU backend (Hygon DTK/HIP) -# ========================= -elseif(USE_DCU) - add_compile_definitions(USE_DCU=1) - - set(DCU_PATH "$ENV{DTK_PATH}" CACHE PATH "Hygon DTK installation root") - set(DCU_ARCH "" CACHE STRING "Optional HIP offload architecture reported by rocminfo") - if(NOT DCU_PATH) - set(DCU_PATH /opt/dtk) - endif() - - find_program(HIPCC_EXECUTABLE hipcc - HINTS "${DCU_PATH}/bin" "${DCU_PATH}/llvm/bin" /opt/rocm/bin - REQUIRED) - if(NOT CMAKE_CXX_COMPILER MATCHES "hipcc") - message(WARNING - "DCU kernels must be compiled by hipcc. Reconfigure with " - "-DCMAKE_CXX_COMPILER=${HIPCC_EXECUTABLE}") - endif() - - include_directories("${DCU_PATH}/include") - link_directories("${DCU_PATH}/lib" "${DCU_PATH}/lib64") - - find_library(DCU_RUNTIME_LIB NAMES amdhip64 hip_hcc - HINTS "${DCU_PATH}/lib" "${DCU_PATH}/lib64" /opt/rocm/lib /opt/rocm/lib64 - REQUIRED) - find_library(DCU_BLAS_LIB NAMES hipblas - HINTS "${DCU_PATH}/lib" "${DCU_PATH}/lib64" /opt/rocm/lib /opt/rocm/lib64 - REQUIRED) - - file(GLOB_RECURSE DCU_KERNELS ${PROJECT_SOURCE_DIR}/infini_train/src/kernels/dcu/*.hip) - set_source_files_properties(${DCU_KERNELS} PROPERTIES - LANGUAGE CXX - COMPILE_OPTIONS "-x;hip" - ) - add_library(infini_train_dcu_kernels STATIC ${DCU_KERNELS}) - if(DCU_ARCH) - target_compile_options(infini_train_dcu_kernels PRIVATE "--offload-arch=${DCU_ARCH}") - endif() - target_link_libraries(infini_train_dcu_kernels glog ${DCU_RUNTIME_LIB} ${DCU_BLAS_LIB}) - - add_library(infini_train STATIC ${SRC}) - target_link_libraries(infini_train glog gflags infini_train_cpu_kernels infini_train_dcu_kernels) - - if(USE_RCCL) - message(STATUS "Add USE_RCCL under DCU backend") - find_library(DCU_COMM_LIB NAMES rccl nccl - HINTS - "${DCU_PATH}/lib" - "${DCU_PATH}/lib64" - "${DCU_PATH}/rccl/lib" - "${DCU_PATH}/cuda/cuda-12/targets/x86_64-linux/lib" - "${DCU_PATH}/cuda/targets/x86_64-linux/lib" - /opt/rocm/lib - /opt/rocm/lib64 - REQUIRED) - target_compile_definitions(infini_train PRIVATE USE_RCCL=1) - target_link_libraries(infini_train ${DCU_COMM_LIB}) - endif() - - if(USE_MPI) - target_link_libraries(infini_train ${MPI_LIBS}) - endif() - -# ========================= -# MACA backend (MetaX) -# ========================= -elseif(USE_MACA) - add_compile_definitions(USE_MACA=1) - - # ---- configure MACA SDK paths ---- - # Typical: /opt/maca (can be overridden by -DMACA_PATH=...) - set(MACA_PATH $ENV{MACA_PATH}) - set(CMAKE_C_COMPILER ${MACA_PATH}/mxgpu_llvm/bin/mxcc) - set(CMAKE_CXX_COMPILER ${MACA_PATH}/mxgpu_llvm/bin/mxcc) - - include_directories("${MACA_PATH}/include") - link_directories("${MACA_PATH}/lib") - - # Libraries: mcruntime / mcdnn / mcblas - find_library(MACA_RUNTIME_LIB NAMES mcruntime HINTS "${MACA_PATH}/lib" REQUIRED) - find_library(MACA_DNN_LIB NAMES mcdnn HINTS "${MACA_PATH}/lib" REQUIRED) - find_library(MACA_BLAS_LIB NAMES mcblas HINTS "${MACA_PATH}/lib" REQUIRED) - - file(GLOB_RECURSE MACA_KERNELS ${PROJECT_SOURCE_DIR}/infini_train/src/kernels/maca/*.maca) - set_source_files_properties(${MACA_KERNELS} PROPERTIES - LANGUAGE CXX - COMPILE_OPTIONS "-x;maca" - ) - add_library(infini_train_maca_kernels STATIC ${MACA_KERNELS}) - target_link_libraries(infini_train_maca_kernels glog ${MACA_RUNTIME_LIB} ${MACA_DNN_LIB} ${MACA_BLAS_LIB}) - - add_library(infini_train STATIC ${SRC}) - target_link_libraries(infini_train glog gflags infini_train_cpu_kernels infini_train_maca_kernels) - - if (USE_MCCL) - message(STATUS "Add USE_MCCL under MACA backend, use MCCL (mccl)") - find_library(MACA_COMM_LIB NAMES mccl HINTS "${MACA_PATH}/lib" REQUIRED) - add_compile_definitions(USE_MCCL=1) - target_link_libraries(infini_train ${MACA_COMM_LIB}) - endif() - - if (USE_MPI) - target_link_libraries(infini_train ${MPI_LIBS} Threads::Threads) - - # 有些 MPI 还需要额外 link flags（比如 -Wl,...），也一并带上 - if (MPI_CXX_LINK_FLAGS) - set_target_properties(infini_train PROPERTIES - LINK_FLAGS "${MPI_CXX_LINK_FLAGS}" - ) - endif() - endif() - -# ========================= -# CPU-only backend -# ========================= -else() - add_library(infini_train STATIC ${SRC}) - target_link_libraries(infini_train glog gflags infini_train_cpu_kernels) ->>>>>>> theirs +endif() + +if(USE_DCU) + target_link_libraries(infini_train PUBLIC infini_train_dcu_kernels ${DCU_RUNTIME_LIB} ${DCU_BLAS_LIB}) + + if(USE_RCCL) + message(STATUS "Add USE_RCCL, use RCCL with DCU") + find_library(DCU_COMM_LIB NAMES rccl nccl + HINTS + "${DCU_PATH}/lib" + "${DCU_PATH}/lib64" + "${DCU_PATH}/rccl/lib" + /opt/rocm/lib + /opt/rocm/lib64 + REQUIRED) + add_compile_definitions(USE_RCCL=1) + target_link_libraries(infini_train PUBLIC ${DCU_COMM_LIB}) + endif() endif() # ------------------------------------------------------------------------------ @@ -311,18 +237,6 @@ function(link_infini_train_exe target_name) "-Wl,--no-whole-archive" "-Wl,--end-group" ) -<<<<<<< ours -======= - elseif(USE_MACA) - target_link_libraries(${target_name} PRIVATE - "-Wl,--start-group" - "-Wl,--whole-archive" - infini_train - infini_train_cpu_kernels - infini_train_maca_kernels - "-Wl,--no-whole-archive" - "-Wl,--end-group" - ) elseif(USE_DCU) target_link_libraries(${target_name} PRIVATE "-Wl,--start-group" @@ -333,7 +247,6 @@ function(link_infini_train_exe target_name) "-Wl,--no-whole-archive" "-Wl,--end-group" ) ->>>>>>> theirs else() target_link_libraries(${target_name} PRIVATE "-Wl,--start-group" diff --git a/README.md b/README.md index f9b70224..1a51d1a8 100644 --- a/README.md +++ b/README.md @@ -35,11 +35,8 @@ A from-scratch C++ training framework for large-scale models with multi-dimensio mkdir build cd build cmake .. -DUSE_CUDA=ON -DUSE_NCCL=ON -make -j -``` - -For Hygon BW1000 / DCU builds and validation, see -[`docs/dcu_bw1000.md`](docs/dcu_bw1000.md). +make -j +``` Build Options: @@ -183,7 +180,6 @@ Multiple parallelism strategies (DDP, TP, SP, PP) can be freely combined to scal Added Autocast, multi-dimensional distributed parallelism (DDP, TP, SP, PP with GPipe / 1F1B / vPP), multi-node training, `no_grad` mode, -<<<<<<< ours and communication–computation overlap with bucketed gradient synchronization. - **2026/06/08** — InfiniTrain **v0.6.0** @@ -212,7 +208,4 @@ Multiple parallelism strategies (DDP, TP, SP, PP) can be freely combined to scal optimizations. Integrated a CTest + GTest based testing infrastructure to strengthen the - framework's automated test workflow. -======= - and communication–computation overlap with bucketed gradient synchronization. ->>>>>>> theirs + framework's automated test workflow. \ No newline at end of file diff --git a/example/gpt2/main.cc b/example/gpt2/main.cc index cce76eb2..b6dcebac 100644 --- a/example/gpt2/main.cc +++ b/example/gpt2/main.cc @@ -75,7 +75,7 @@ DEFINE_uint32(sample_every, 0, "how often to sample from the model?"); // debugging DEFINE_bool(overfit_single_batch, true, "overfit just one batch of data"); // memory management -DEFINE_string(device, "cuda", "device type (cpu/cuda/maca/dcu), useless if using parallel training mode"); +DEFINE_string(device, "cuda", "device type (cpu/cuda/dcu), useless if using parallel training mode"); // parallel DEFINE_int32( nthread_per_process, 1, @@ -114,11 +114,7 @@ const std::unordered_set kSupportedModels = {"gpt2", "gpt2-medium", "gpt2-large", "gpt2-xl", "d12", "d24", "d36", "d48"}; constexpr char kDeviceCPU[] = "cpu"; constexpr char kDeviceCUDA[] = "cuda"; -<<<<<<< ours -======= -constexpr char kDeviceMACA[] = "maca"; constexpr char kDeviceDCU[] = "dcu"; ->>>>>>> theirs constexpr char kDtypeFP32[] = "float32"; constexpr char kDtypeBF16[] = "bfloat16"; const std::unordered_set kSupportedLRDecayStyles @@ -135,17 +131,12 @@ const std::unordered_map kModelToConfigs = { } // namespace DEFINE_validator(model, [](const char *, const std::string &value) { return kSupportedModels.contains(value); }); -<<<<<<< ours -DEFINE_validator(device, - [](const char *, const std::string &value) { return value == kDeviceCPU || value == kDeviceCUDA; }); +DEFINE_validator(device, [](const char *, const std::string &value) { + return value == kDeviceCPU || value == kDeviceCUDA || value == kDeviceDCU; +}); DEFINE_validator(zero_stage, [](const char *, int32_t value) { return value >= 0 && value <= 3; }); DEFINE_validator(lr_decay_style, [](const char *, const std::string &value) { return kSupportedLRDecayStyles.contains(value); }); -======= -DEFINE_validator(device, [](const char *, const std::string &value) { - return value == kDeviceCPU || value == kDeviceCUDA || value == kDeviceMACA || value == kDeviceDCU; -}); ->>>>>>> theirs void Train(const nn::parallel::Rank &rank) { using namespace nn::parallel; @@ -192,18 +183,8 @@ void Train(const nn::parallel::Rank &rank) { const ProcessGroup *pp_pg = nullptr; if (rank.IsParallel()) { -<<<<<<< ours - device = Device(Device::DeviceType::kCUDA, rank.thread_rank()); -======= - auto parallel_device_type = Device::DeviceType::kCUDA; - if (FLAGS_device == kDeviceMACA) { - parallel_device_type = Device::DeviceType::kMACA; - } else if (FLAGS_device == kDeviceDCU) { - parallel_device_type = Device::DeviceType::kDCU; - } + auto parallel_device_type = FLAGS_device == kDeviceDCU ? Device::DeviceType::kDCU : Device::DeviceType::kCUDA; device = Device(parallel_device_type, rank.thread_rank()); - ->>>>>>> theirs auto *pg_factory = ProcessGroupFactory::Instance(device.type()); if (ddp_world_size > 1) { @@ -227,20 +208,12 @@ void Train(const nn::parallel::Rank &rank) { nn::parallel::pp_rank = pp_rank; } + } else if (FLAGS_device == kDeviceCPU) { + device = Device(); + } else if (FLAGS_device == kDeviceDCU) { + device = Device(Device::DeviceType::kDCU, 0); } else { -<<<<<<< ours - device = FLAGS_device == kDeviceCPU ? Device() : Device(Device::DeviceType::kCUDA, 0); -======= - if (FLAGS_device == kDeviceCPU) { - device = Device(); - } else if (FLAGS_device == kDeviceMACA) { - device = Device(Device::DeviceType::kMACA, 0); - } else if (FLAGS_device == kDeviceDCU) { - device = Device(Device::DeviceType::kDCU, 0); - } else { - device = Device(Device::DeviceType::kCUDA, 0); - } ->>>>>>> theirs + device = Device(Device::DeviceType::kCUDA, 0); } // calculate gradient accumulation from the desired total batch size and the current run configuration diff --git a/example/llama3/main.cc b/example/llama3/main.cc index e41703cd..d150024b 100644 --- a/example/llama3/main.cc +++ b/example/llama3/main.cc @@ -74,7 +74,7 @@ DEFINE_uint32(sample_every, 0, "how often to sample from the model?"); // debugging DEFINE_bool(overfit_single_batch, true, "overfit just one batch of data"); // memory management -DEFINE_string(device, "cuda", "device type (cpu/cuda/maca/dcu), useless if using parallel training mode"); +DEFINE_string(device, "cuda", "device type (cpu/cuda/dcu), useless if using parallel training mode"); // parallel DEFINE_int32( nthread_per_process, 1, @@ -110,11 +110,7 @@ namespace { const std::unordered_set kSupportedModels = {"llama3"}; constexpr char kDeviceCPU[] = "cpu"; constexpr char kDeviceCUDA[] = "cuda"; -<<<<<<< ours -======= -constexpr char kDeviceMACA[] = "maca"; constexpr char kDeviceDCU[] = "dcu"; ->>>>>>> theirs constexpr char kDtypeFP32[] = "float32"; constexpr char kDtypeBF16[] = "bfloat16"; const std::unordered_set kSupportedLRDecayStyles @@ -122,17 +118,12 @@ const std::unordered_set kSupportedLRDecayStyles } // namespace DEFINE_validator(model, [](const char *, const std::string &value) { return kSupportedModels.contains(value); }); -<<<<<<< ours -DEFINE_validator(device, - [](const char *, const std::string &value) { return value == kDeviceCPU || value == kDeviceCUDA; }); +DEFINE_validator(device, [](const char *, const std::string &value) { + return value == kDeviceCPU || value == kDeviceCUDA || value == kDeviceDCU; +}); DEFINE_validator(zero_stage, [](const char *, int32_t value) { return value >= 0 && value <= 3; }); DEFINE_validator(lr_decay_style, [](const char *, const std::string &value) { return kSupportedLRDecayStyles.contains(value); }); -======= -DEFINE_validator(device, [](const char *, const std::string &value) { - return value == kDeviceCPU || value == kDeviceCUDA || value == kDeviceMACA || value == kDeviceDCU; -}); ->>>>>>> theirs void Train(const nn::parallel::Rank &rank) { using namespace nn::parallel; @@ -178,18 +169,8 @@ void Train(const nn::parallel::Rank &rank) { const ProcessGroup *pp_pg = nullptr; if (rank.IsParallel()) { -<<<<<<< ours - device = Device(Device::DeviceType::kCUDA, rank.thread_rank()); -======= - auto parallel_device_type = Device::DeviceType::kCUDA; - if (FLAGS_device == kDeviceMACA) { - parallel_device_type = Device::DeviceType::kMACA; - } else if (FLAGS_device == kDeviceDCU) { - parallel_device_type = Device::DeviceType::kDCU; - } + auto parallel_device_type = FLAGS_device == kDeviceDCU ? Device::DeviceType::kDCU : Device::DeviceType::kCUDA; device = Device(parallel_device_type, rank.thread_rank()); - ->>>>>>> theirs auto *pg_factory = ProcessGroupFactory::Instance(device.type()); if (ddp_world_size > 1) { @@ -213,20 +194,12 @@ void Train(const nn::parallel::Rank &rank) { nn::parallel::pp_rank = pp_rank; } + } else if (FLAGS_device == kDeviceCPU) { + device = Device(); + } else if (FLAGS_device == kDeviceDCU) { + device = Device(Device::DeviceType::kDCU, 0); } else { -<<<<<<< ours - device = FLAGS_device == kDeviceCPU ? Device() : Device(Device::DeviceType::kCUDA, 0); -======= - if (FLAGS_device == kDeviceCPU) { - device = Device(); - } else if (FLAGS_device == kDeviceMACA) { - device = Device(Device::DeviceType::kMACA, 0); - } else if (FLAGS_device == kDeviceDCU) { - device = Device(Device::DeviceType::kDCU, 0); - } else { - device = Device(Device::DeviceType::kCUDA, 0); - } ->>>>>>> theirs + device = Device(Device::DeviceType::kCUDA, 0); } // calculate gradient accumulation from the desired total batch size and the current run configuration diff --git a/infini_train/include/device.h b/infini_train/include/device.h index 8942d9e5..11d34142 100644 --- a/infini_train/include/device.h +++ b/infini_train/include/device.h @@ -13,13 +13,8 @@ class Device { enum class DeviceType : int8_t { kCPU = 0, kCUDA = 1, -<<<<<<< ours - kCount = 2, -======= - kMACA = 2, - kDCU = 3, - kCount = 4, ->>>>>>> theirs + kDCU = 2, + kCount = 3, kInvalid = -1, }; @@ -36,11 +31,7 @@ class Device { bool IsCPU() const; bool IsCUDA() const; -<<<<<<< ours -======= - bool IsMACA() const; bool IsDCU() const; ->>>>>>> theirs std::string ToString() const; diff --git a/infini_train/src/device.cc b/infini_train/src/device.cc index 296969f3..a233e798 100644 --- a/infini_train/src/device.cc +++ b/infini_train/src/device.cc @@ -25,13 +25,6 @@ bool Device::IsCPU() const { return type_ == DeviceType::kCPU; } bool Device::IsCUDA() const { return type_ == DeviceType::kCUDA; } -<<<<<<< ours -std::string Device::ToString() const { - std::ostringstream oss; - oss << std::format("Device({}, {})", type_ == DeviceType::kCPU ? "CPU" : "CUDA", index_); -======= -bool Device::IsMACA() const { return type_ == DeviceType::kMACA; } - bool Device::IsDCU() const { return type_ == DeviceType::kDCU; } std::string Device::ToString() const { @@ -43,9 +36,6 @@ std::string Device::ToString() const { case DeviceType::kCUDA: type_str = "CUDA"; break; - case DeviceType::kMACA: - type_str = "MACA"; - break; case DeviceType::kDCU: type_str = "DCU"; break; @@ -54,7 +44,6 @@ std::string Device::ToString() const { } std::ostringstream oss; oss << "Device(" << type_str << ", " << static_cast(index_) << ")"; ->>>>>>> theirs return oss.str(); } From c803e357cbfe104af7ca6e30e956830cbf4c8a17 Mon Sep 17 00:00:00 2001 From: cy <2833839179@qq.com> Date: Thu, 6 Aug 2026 09:13:32 +0800 Subject: [PATCH 4/4] fix: make DCU backend compatible with master --- infini_train/include/autograd/linear.h | 6 + infini_train/include/dtype_dispatch.h | 1 + infini_train/src/kernels/dcu/linear.hip | 148 ++++++++++++++++++++++++ infini_train/src/nn/modules/module.cc | 32 +++-- 4 files changed, 180 insertions(+), 7 deletions(-) diff --git a/infini_train/include/autograd/linear.h b/infini_train/include/autograd/linear.h index cebed3b2..21d107b9 100644 --- a/infini_train/include/autograd/linear.h +++ b/infini_train/include/autograd/linear.h @@ -12,6 +12,12 @@ class Tensor; namespace infini_train::autograd { +struct LinearGradFlags { + bool input = false; + bool weight = false; + bool bias = false; +}; + class Linear : public Function { public: static constexpr char kType[] = "LinearFunction"; diff --git a/infini_train/include/dtype_dispatch.h b/infini_train/include/dtype_dispatch.h index 1847fcc5..302a4717 100644 --- a/infini_train/include/dtype_dispatch.h +++ b/infini_train/include/dtype_dispatch.h @@ -184,6 +184,7 @@ namespace infini_train { #define INFINI_SIGNED_INTEGRAL_TYPES DataType::kINT8, DataType::kINT16, DataType::kINT32, DataType::kINT64 #define INFINI_UNSIGNED_INTEGRAL_TYPES DataType::kUINT8, DataType::kUINT16, DataType::kUINT32, DataType::kUINT64 #define INFINI_ALL_INTEGRAL_TYPES INFINI_SIGNED_INTEGRAL_TYPES, INFINI_UNSIGNED_INTEGRAL_TYPES +#define INFINI_ALL_TYPES INFINI_ALL_FLOATING_TYPES, INFINI_ALL_INTEGRAL_TYPES #define INFINI_ALL_NUMERIC_TYPES INFINI_ALL_FLOATING_TYPES, INFINI_ALL_INTEGRAL_TYPES #define INFINI_8_BIT_TYPES DataType::kINT8, DataType::kUINT8 #define INFINI_16_BIT_TYPES DataType::kINT16, DataType::kUINT16, DataType::kFLOAT16, DataType::kBFLOAT16 diff --git a/infini_train/src/kernels/dcu/linear.hip b/infini_train/src/kernels/dcu/linear.hip index f72d41cb..07d61391 100644 --- a/infini_train/src/kernels/dcu/linear.hip +++ b/infini_train/src/kernels/dcu/linear.hip @@ -188,6 +188,123 @@ MatmulBackward(const std::shared_ptr &input, const std::shared_ptr MatmulBackwardInput(const std::shared_ptr &other, + const std::shared_ptr &grad_output, + const std::vector &input_dims) { + /* + grad_input[*, m, k] = grad_output[*, m, n] * other[*, k, n]^T + */ + const auto &other_dims = other->Dims(); + const auto &grad_output_dims = grad_output->Dims(); + + CHECK_GE(other_dims.size(), 2); + CHECK_EQ(other_dims.size(), grad_output_dims.size()); + + const int64_t m = grad_output_dims[grad_output_dims.size() - 2]; + const int64_t k = other_dims[other_dims.size() - 2]; + const int64_t n = other_dims[other_dims.size() - 1]; + CHECK_EQ(n, grad_output_dims[grad_output_dims.size() - 1]); + + const int64_t bs + = std::accumulate(grad_output_dims.rbegin() + 2, grad_output_dims.rend(), 1, std::multiplies{}); + for (int64_t i = 0; i < static_cast(grad_output_dims.size()) - 2; ++i) { + CHECK_EQ(grad_output_dims[i], other_dims[i]) << "Batch dims must match"; + } + + auto compute_dtype = other->Dtype(); + auto grad_output_dtype = grad_output->Dtype(); + auto grad_output_promoted + = grad_output_dtype == compute_dtype ? grad_output : std::make_shared(grad_output->To(compute_dtype)); + auto output_dtype = (compute_dtype == DataType::kBFLOAT16) ? DataType::kFLOAT32 : compute_dtype; + auto grad_input = std::make_shared(input_dims, output_dtype, grad_output->GetDevice()); + + auto device = grad_output->GetDevice(); + hipblasHandle_t handle = dynamic_cast( + infini_train::core::GetDeviceGuardImpl(device.type())->GetBlasHandle(device)) + ->hipblas_handle(); + + const float alpha = 1.0f, beta = 0.0f; + const int lda = n, ldb = n, ldc = k; + const int64_t stride_a = k * n; + const int64_t stride_b = n * m; + const int64_t stride_c = m * k; + switch (compute_dtype) { + DISPATCH_CASE(WRAP(HIPBLAS_CHECK(hipblasGemmStridedBatchedEx( + handle, HIPBLAS_OP_T, HIPBLAS_OP_N, k, m, n, &alpha, other->DataPtr(), HIPBLAS_R_32F, lda, + stride_a, grad_output_promoted->DataPtr(), HIPBLAS_R_32F, ldb, stride_b, &beta, + grad_input->DataPtr(), HIPBLAS_R_32F, ldc, stride_c, bs, HIPBLAS_R_32F, + HIPBLAS_GEMM_DEFAULT));), + DataType::kFLOAT32) + DISPATCH_CASE(WRAP(HIPBLAS_CHECK(hipblasGemmStridedBatchedEx( + handle, HIPBLAS_OP_T, HIPBLAS_OP_N, k, m, n, &alpha, other->DataPtr(), HIPBLAS_R_16B, lda, + stride_a, grad_output_promoted->DataPtr(), HIPBLAS_R_16B, ldb, stride_b, &beta, + grad_input->DataPtr(), HIPBLAS_R_32F, ldc, stride_c, bs, HIPBLAS_R_32F, + HIPBLAS_GEMM_DEFAULT));), + DataType::kBFLOAT16) + } + + return grad_input; +} + +std::shared_ptr MatmulBackwardOther(const std::shared_ptr &input, + const std::shared_ptr &grad_output, + const std::vector &other_dims) { + /* + grad_other[*, k, n] = input[*, m, k]^T * grad_output[*, m, n] + */ + const auto &input_dims = input->Dims(); + const auto &grad_output_dims = grad_output->Dims(); + + CHECK_GE(input_dims.size(), 2); + CHECK_EQ(input_dims.size(), grad_output_dims.size()); + + const int64_t m = input_dims[input_dims.size() - 2]; + const int64_t k = input_dims[input_dims.size() - 1]; + const int64_t n = grad_output_dims[grad_output_dims.size() - 1]; + CHECK_EQ(m, grad_output_dims[grad_output_dims.size() - 2]); + CHECK_EQ(k, other_dims[other_dims.size() - 2]); + + const int64_t bs = std::accumulate(input_dims.rbegin() + 2, input_dims.rend(), 1, std::multiplies{}); + for (int64_t i = 0; i < static_cast(input_dims.size()) - 2; ++i) { + CHECK_EQ(input_dims[i], grad_output_dims[i]) << "Batch dims must match"; + CHECK_EQ(input_dims[i], other_dims[i]) << "Batch dims must match"; + } + + auto compute_dtype = input->Dtype(); + auto grad_output_dtype = grad_output->Dtype(); + auto grad_output_promoted + = grad_output_dtype == compute_dtype ? grad_output : std::make_shared(grad_output->To(compute_dtype)); + auto output_dtype = (compute_dtype == DataType::kBFLOAT16) ? DataType::kFLOAT32 : compute_dtype; + auto grad_other = std::make_shared(other_dims, output_dtype, grad_output->GetDevice()); + + auto device = grad_output->GetDevice(); + hipblasHandle_t handle = dynamic_cast( + infini_train::core::GetDeviceGuardImpl(device.type())->GetBlasHandle(device)) + ->hipblas_handle(); + + const float alpha = 1.0f, beta = 0.0f; + const int lda = n, ldb = k, ldc = n; + const int64_t stride_a = n * m; + const int64_t stride_b = k * m; + const int64_t stride_c = n * k; + switch (compute_dtype) { + DISPATCH_CASE(WRAP(HIPBLAS_CHECK(hipblasGemmStridedBatchedEx( + handle, HIPBLAS_OP_N, HIPBLAS_OP_T, n, k, m, &alpha, grad_output_promoted->DataPtr(), + HIPBLAS_R_32F, lda, stride_a, input->DataPtr(), HIPBLAS_R_32F, ldb, stride_b, &beta, + grad_other->DataPtr(), HIPBLAS_R_32F, ldc, stride_c, bs, HIPBLAS_R_32F, + HIPBLAS_GEMM_DEFAULT));), + DataType::kFLOAT32) + DISPATCH_CASE(WRAP(HIPBLAS_CHECK(hipblasGemmStridedBatchedEx( + handle, HIPBLAS_OP_N, HIPBLAS_OP_T, n, k, m, &alpha, grad_output_promoted->DataPtr(), + HIPBLAS_R_16B, lda, stride_a, input->DataPtr(), HIPBLAS_R_16B, ldb, stride_b, &beta, + grad_other->DataPtr(), HIPBLAS_R_32F, ldc, stride_c, bs, HIPBLAS_R_32F, + HIPBLAS_GEMM_DEFAULT));), + DataType::kBFLOAT16) + } + + return grad_other; +} + template __global__ void BiasCopyKernel(T *output, const T *bias, int bs, int out_features) { int idx = blockIdx.x * blockDim.x + threadIdx.x; if (idx >= bs * out_features) { @@ -491,6 +608,32 @@ LinearBackward(const std::shared_ptr &input, const std::shared_ptr LinearBackwardInput(const std::shared_ptr &weight, + const std::shared_ptr &grad_output, bool transpose, + int64_t in_features, int64_t out_features, + const std::vector &input_dims) { + auto [grad_input, grad_weight, grad_bias] + = LinearBackward(nullptr, weight, transpose, in_features, out_features, input_dims, grad_output, false, + infini_train::autograd::LinearGradFlags{.input = true}); + return grad_input; +} + +std::shared_ptr LinearBackwardWeight(const std::shared_ptr &input, + const std::shared_ptr &grad_output, bool transpose, + int64_t in_features, int64_t out_features) { + auto [grad_input, grad_weight, grad_bias] + = LinearBackward(input, nullptr, transpose, in_features, out_features, input->Dims(), grad_output, false, + infini_train::autograd::LinearGradFlags{.weight = true}); + return grad_weight; +} + +std::shared_ptr LinearBackwardBias(const std::shared_ptr &grad_output, int64_t out_features) { + auto [grad_input, grad_weight, grad_bias] + = LinearBackward(nullptr, nullptr, true, 0, out_features, grad_output->Dims(), grad_output, true, + infini_train::autograd::LinearGradFlags{.bias = true}); + return grad_bias; +} } // namespace infini_train::kernels::dcu #define REGISTER_HIP_LINEAR_KERNEL(kernel_name) \ @@ -498,7 +641,12 @@ LinearBackward(const std::shared_ptr &input, const std::shared_ptr #include +#include #include #include #include @@ -20,6 +21,26 @@ #endif namespace infini_train::nn { +namespace { + +std::string LoadStateDictMissingKeyMessage(const std::string &name) { return "Missing key: " + name; } + +std::string LoadStateDictShapeMismatchMessage(const std::string &name, const std::vector &expected, + const std::vector &actual) { + std::ostringstream oss; + oss << "Shape mismatch for '" << name << "': expected=" << infini_train::utils::DimsToString(expected) + << ", got=" << infini_train::utils::DimsToString(actual); + return oss.str(); +} + +std::string LoadStateDictDtypeMismatchMessage(const std::string &name, DataType expected, DataType actual) { + std::ostringstream oss; + oss << "Dtype mismatch for '" << name << "': expected=" << kDataTypeToDesc.at(expected) + << ", got=" << kDataTypeToDesc.at(actual); + return oss.str(); +} + +} // namespace Module::Module() : Module(kUndefinedType) {} @@ -156,24 +177,21 @@ void Module::LoadStateDict(const std::unordered_mapDims() != src->Dims()) { - error_msgs.push_back(std::format("Shape mismatch for '{}': expected={}, got={}", name, - infini_train::utils::DimsToString(dst->Dims()), - infini_train::utils::DimsToString(src->Dims()))); + error_msgs.push_back(LoadStateDictShapeMismatchMessage(name, dst->Dims(), src->Dims())); } if (dst->Dtype() != src->Dtype()) { - error_msgs.push_back(std::format("Dtype mismatch for '{}': expected={}, got={}", name, - kDataTypeToDesc.at(dst->Dtype()), kDataTypeToDesc.at(src->Dtype()))); + error_msgs.push_back(LoadStateDictDtypeMismatchMessage(name, dst->Dtype(), src->Dtype())); } } for (const auto &[name, src] : state_dict) { if (!visited_keys.contains(name)) { - LOG(WARNING) << std::format("Unexpected key in state_dict: {}", name); + LOG(WARNING) << "Unexpected key in state_dict: " << name; } }