diff --git a/.clang-format b/.clang-format new file mode 100644 index 0000000..e26a42f --- /dev/null +++ b/.clang-format @@ -0,0 +1,30 @@ +--- +BasedOnStyle: LLVM +IndentWidth: 4 +AccessModifierOffset: -4 +AlignOperands: AlignAfterOperator +BreakBeforeBinaryOperators: All +ColumnLimit: 120 +AllowShortBlocksOnASingleLine: Always +AllowShortLoopsOnASingleLine: true +InsertBraces: true +BreakBeforeBraces: Custom +BraceWrapping: + AfterCaseLabel: false + AfterClass: false + AfterControlStatement: Never + AfterEnum: false + AfterFunction: false + AfterNamespace: false + AfterObjCDeclaration: false + AfterStruct: false + AfterUnion: false + AfterExternBlock: false + BeforeCatch: false + BeforeElse: false + BeforeLambdaBody: false + BeforeWhile: false + IndentBraces: false + SplitEmptyFunction: true + SplitEmptyRecord: true + SplitEmptyNamespace: true diff --git a/.github/workflows/format-check.yaml b/.github/workflows/format-check.yaml new file mode 100644 index 0000000..1fc2def --- /dev/null +++ b/.github/workflows/format-check.yaml @@ -0,0 +1,27 @@ +name: Format Check + +on: + pull_request: + push: + paths-ignore: + - '**.md' + - 'LICENSE' + +jobs: + format-check: + name: Check Code Format + runs-on: ubuntu-latest + steps: + + - name: Checkout code + uses: actions/checkout@v4 + + - name: Install Python dependencies + run: | + python3 -m pip install --upgrade pip + pip install black + + - name: Run format check + run: | + python3 scripts/format.py --path backends --check + diff --git a/.gitignore b/.gitignore new file mode 100644 index 0000000..f9047c8 --- /dev/null +++ b/.gitignore @@ -0,0 +1,9 @@ +/build/ +/cmake-build-*/ +/.cache/ +/.vscode/ +/compile_commands.json + +*.log +*.report.rank* +*.records.log.rank* diff --git a/.gitmodules b/.gitmodules new file mode 100644 index 0000000..719b3bc --- /dev/null +++ b/.gitmodules @@ -0,0 +1,3 @@ +[submodule "third_party/InfiniTrain"] + path = third_party/InfiniTrain + url = https://github.com/InfiniTensor/InfiniTrain.git diff --git a/CMakeLists.txt b/CMakeLists.txt new file mode 100644 index 0000000..86e1774 --- /dev/null +++ b/CMakeLists.txt @@ -0,0 +1,71 @@ +cmake_minimum_required(VERSION 3.28) + +set(INFINITRAIN_BACKEND "" CACHE STRING + "Accelerator provider to build from backends/") +if(INFINITRAIN_BACKEND STREQUAL "") + message(FATAL_ERROR + "INFINITRAIN_BACKEND is required. Configure with " + "-DINFINITRAIN_BACKEND=.") +endif() +string(TOLOWER "${INFINITRAIN_BACKEND}" INFINITRAIN_BACKEND) +if(NOT INFINITRAIN_BACKEND MATCHES "^[a-z0-9_]+$") + message(FATAL_ERROR + "INFINITRAIN_BACKEND must name a directory under backends/: " + "${INFINITRAIN_BACKEND}") +endif() + +set(INFINITRAIN_BACKEND_DIR + "${CMAKE_CURRENT_SOURCE_DIR}/backends/${INFINITRAIN_BACKEND}") +set(INFINITRAIN_BACKEND_PRE_PROJECT + "${INFINITRAIN_BACKEND_DIR}/cmake/pre_project.cmake") +if(NOT EXISTS "${INFINITRAIN_BACKEND_PRE_PROJECT}") + message(FATAL_ERROR + "Unknown or incomplete InfiniTrain backend '${INFINITRAIN_BACKEND}': " + "${INFINITRAIN_BACKEND_PRE_PROJECT} was not found.") +endif() + +if(USE_CUDA) + message(FATAL_ERROR + "USE_CUDA=ON is incompatible with the '${INFINITRAIN_BACKEND}' " + "PrivateUse1 provider. Configure with -DUSE_CUDA=OFF or use a separate " + "build directory.") +endif() + +# A provider owns compiler and dependency-probe setup that must happen before +# project(). This keeps vendor SDK assumptions out of the repository root. +include("${INFINITRAIN_BACKEND_PRE_PROJECT}") + +project(InfiniTrainBackends VERSION 0.1.0 LANGUAGES CXX) + +option(INFINITRAIN_BACKENDS_BUILD_EXAMPLES + "Build backend-enabled InfiniTrain examples" ${PROJECT_IS_TOP_LEVEL}) + +set(CMAKE_CXX_STANDARD 20) +set(CMAKE_CXX_STANDARD_REQUIRED ON) +set(CMAKE_CXX_EXTENSIONS OFF) + +set(INFINITRAIN_SOURCE_DIR + "${CMAKE_CURRENT_SOURCE_DIR}/third_party/InfiniTrain" + CACHE PATH "InfiniTrain source tree (normally the pinned submodule)") +if(NOT EXISTS "${INFINITRAIN_SOURCE_DIR}/CMakeLists.txt") + message(FATAL_ERROR + "InfiniTrain was not found at ${INFINITRAIN_SOURCE_DIR}. " + "Run git submodule update --init --recursive or set INFINITRAIN_SOURCE_DIR.") +endif() + +add_subdirectory("${INFINITRAIN_SOURCE_DIR}" "${CMAKE_CURRENT_BINARY_DIR}/InfiniTrain") +if(INFINITRAIN_BACKENDS_BUILD_EXAMPLES AND NOT TARGET infini_run) + # InfiniTrain omits top-level tools when embedded, but multi-process model + # runs still require its launcher next to the provider-enabled examples. + add_subdirectory( + "${INFINITRAIN_SOURCE_DIR}/tools/infini_run" + "${CMAKE_CURRENT_BINARY_DIR}/InfiniTrain/tools/infini_run") + set_target_properties(infini_run PROPERTIES + RUNTIME_OUTPUT_DIRECTORY "${CMAKE_BINARY_DIR}") +endif() +if(BUILD_TEST) + # InfiniTrain registers the shared suites in its subdirectory. Enable CTest + # at this repository root before the provider instantiates PrivateUse1. + enable_testing() +endif() +add_subdirectory("backends/${INFINITRAIN_BACKEND}") diff --git a/README.md b/README.md index ff4aab8..db8e780 100644 --- a/README.md +++ b/README.md @@ -1 +1,231 @@ -# InfiniTrain-Backends \ No newline at end of file +# InfiniTrain Backends + +[![Issues](https://img.shields.io/github/issues/InfiniTensor/InfiniTrain-Backends)]( +https://github.com/InfiniTensor/InfiniTrain-Backends/issues +) +[![PR](https://img.shields.io/github/issues-pr/InfiniTensor/InfiniTrain-Backends)]( +https://github.com/InfiniTensor/InfiniTrain-Backends/pulls +) +[![License](https://img.shields.io/github/license/InfiniTensor/InfiniTrain-Backends)]( +https://github.com/InfiniTensor/InfiniTrain-Backends/blob/master/LICENSE +) + +InfiniTrain Backends provides out-of-tree accelerator backends for +[InfiniTrain](https://github.com/InfiniTensor/InfiniTrain). It keeps +vendor-specific SDK integration, runtime support, collective communication, +and kernels outside the framework core while implementing InfiniTrain's +`PrivateUse1` backend interfaces. + +Each provider is isolated under `backends/` and normally consumes a +pinned InfiniTrain submodule commit. One build tree selects one provider, and a +process may register at most one provider for `DeviceType::kPrivateUse1`. + +## Supported Backends + +| Backend | Runtime | Collectives | Model examples | +| ------- | ------- | ----------- | -------------- | +| MACA | MACA | MCCL (optional) | GPT-2, LLaMA 3, Mixtral | + +The model sources normally come from the pinned InfiniTrain submodule. This +repository provides the provider-specific runtime, kernels, collective +implementation, and build integration needed to run them on MACA. + +## Requirements + +- Linux +- CMake 3.28 or newer +- Git with submodule support +- A compatible MACA SDK with a C++20-capable `mxgpu_llvm/bin/mxcc` compiler, + the MACA runtime, MCDNN, and MCBLAS +- MCCL when distributed collectives are enabled +- `jq` when using the automated model test runner + +## Quick Start + +Initialize the pinned InfiniTrain submodule, select the MACA SDK, and build the +project: + +```bash +git submodule update --init --recursive +export MACA_PATH=/opt/maca + +mkdir build +cd build +cmake .. \ + -DINFINITRAIN_BACKEND=maca \ + -DINFINITRAIN_MACA_WITH_MCCL=ON \ + -DBUILD_TEST=ON +make -j +``` + +Top-level builds enable the backend examples by default. The example +executables are written to `build`: + +```bash +./gpt2 --help +./llama3 --help +./mixtral --help +``` + +Run the registered MACA accelerator tests with CTest: + +```bash +ctest -L maca --output-on-failure +``` + +## Training + +As in InfiniTrain, each model example is an independent executable. Select the +MACA backend with `--device maca`. For example, a single-node LLaMA 3 training +run can be started from the build directory with: + +```bash +./llama3 \ + --device maca \ + --input_bin [training_data_path] \ + --llmc_filepath [model_path] \ + --num_iteration 10 +``` + +The GPT-2 and Mixtral executables follow the same command-line interface for +their corresponding model and dataset options. Run an executable with `--help` +to inspect all available options. + +The model test matrix reuses InfiniTrain's test runner with MACA-specific +configuration. Update the dataset and checkpoint paths in +`backends/maca/scripts/test_config_maca.json` before running it from the +repository root in a separate shell: + +```bash +backends/maca/scripts/run_models_and_profile.bash --only-run basic +``` + +## Build Options + +| Option | Default | Description | +| ------ | ------- | ----------- | +| `INFINITRAIN_BACKEND` | Required | Provider selected from `backends/` | +| `INFINITRAIN_SOURCE_DIR` | `third_party/InfiniTrain` | InfiniTrain source tree, normally the pinned submodule | +| `INFINITRAIN_BACKENDS_BUILD_EXAMPLES` | `ON` for a top-level build | Build provider-enabled InfiniTrain examples | +| `BUILD_TEST` | `OFF` | Build InfiniTrain's full test set and MACA variants | +| `INFINITRAIN_MACA_WITH_MCCL` | `ON` | Enable MCCL distributed collectives | +| `MACA_PATH` | `$MACA_PATH` | MACA SDK root | + +One build directory may contain only one provider. When working with multiple +providers, use a separate directory for each one so compiler and SDK cache +entries do not leak between them. Starting from the repository root: + +```bash +mkdir build-maca +cd build-maca +cmake .. -DINFINITRAIN_BACKEND=maca +``` + +For development against another InfiniTrain checkout, override the pinned +submodule path explicitly: + +```bash +mkdir build +cd build +cmake .. \ + -DINFINITRAIN_BACKEND=maca \ + -DINFINITRAIN_SOURCE_DIR=/path/to/InfiniTrain +``` + +The selected checkout must implement the PrivateUse1 extension API expected by +this backend. + +To instantiate the shared accelerator tests, configure with `BUILD_TEST=ON`. +PrivateUse1 providers require `USE_CUDA=OFF`; configuration fails rather than +silently overriding an explicit `USE_CUDA=ON`. Their test identity is always +PrivateUse1, independent of the selected provider: + +```bash +cmake .. -DBUILD_TEST=ON +cmake --build . --target test_tensor_maca test_autograd_maca +ctest -L maca --output-on-failure +``` + +The generated binaries are named `test_*_maca`, contain only +`PRIVATEUSE1/*` GTest instances, and carry the `maca`, `accelerator`, and +`hardware` CTest labels. The same build also contains InfiniTrain's CPU, +fake-provider, and CPU-only tests; CUDA remains disabled by the PrivateUse1 +configuration contract. Use `ctest -L cpu` for the upstream CPU tests, or run +`ctest --output-on-failure` without `-L` to execute the complete registered set. + +## Using the MACA Backend + +Applications must register the provider before parsing or constructing a +`maca` device: + +```cpp +#include "infini_train_maca/backend.h" + +infini_train::maca::RegisterBackend(); +auto type = infini_train::Device::ParseType("maca").value(); +infini_train::Device device(type, 0); +``` + +`RegisterBackend()` installs the process-wide PrivateUse1 name, MACA kernels, +device guard, and, when enabled, the MCCL implementation. Registration itself +does not initialize the device runtime. MACA runtime initialization remains +lazy until InfiniTrain first requests the device guard. + +Final executables should link `InfiniTrain::Backend::MACAExecutable` instead of +assembling the static archives themselves: + +```cmake +add_executable(train main.cc) +target_link_libraries(train PRIVATE InfiniTrain::Backend::MACAExecutable) +``` + +This interface target retains InfiniTrain's static kernel-registration objects +and links the MACA provider in the required order. Library targets that do not +produce a final executable may link `InfiniTrain::Backend::MACA`. + +## MACA Runtime Notes + +- `MACA_LAUNCH_BLOCKING`: Unless already set by the user, the provider sets it + to `1` immediately before lazy runtime initialization. Set the variable + before the first device use to override it. +- `MCCL_P2P_DISABLE`: Immediately before runtime initialization, the provider + reads `--tensor_parallel` from the process command line and sets this to `1` + when the value is greater than one. An explicit environment value is + preserved. Data-parallel jobs leave it unset to retain MCCL's P2P fast path. + +## Architecture + +```text +backends// + cmake/ compiler and vendor SDK setup before project() + include/ public provider API + src/backend.cc provider registration entry point + src/common/ provider-internal shared helpers + src/runtime/ device, stream, event, and allocator integration + src/kernels/ provider kernel implementations and registration + src/ccl/ optional collective communication integration + examples/ CMake adapters for upstream InfiniTrain examples + scripts/ provider model-run configuration and wrappers + +third_party/InfiniTrain/ framework source, normally a pinned git submodule +``` + +The root build selects and includes +`backends//cmake/pre_project.cmake` before its first `project()` call. +This lets each provider select its compiler and prepare SDK-specific dependency +probes without adding vendor branches to the root `CMakeLists.txt`. + +The MACA implementation uses InfiniTrain's existing registration mechanisms: +`REGISTER_KERNEL`, `INFINI_TRAIN_REGISTER_DEVICE_GUARD_IMPL`, and +`INFINI_TRAIN_REGISTER_CCL_IMPL`. `RegisterBackend()` explicitly reaches the +runtime and kernel paths, plus the CCL path when MCCL is enabled. Static-library +object extraction is therefore driven by strong symbol references instead of +depending on global initialization order. + +When this repository is embedded with `add_subdirectory()`, the parent project +must select `${MACA_PATH}/mxgpu_llvm/bin/mxcc` before its first `project()` +call. CMake cannot replace a compiler after a language has been enabled. + +## License + +InfiniTrain Backends is released under the [MIT License](LICENSE). diff --git a/backends/maca/CMakeLists.txt b/backends/maca/CMakeLists.txt new file mode 100644 index 0000000..9f65328 --- /dev/null +++ b/backends/maca/CMakeLists.txt @@ -0,0 +1,98 @@ +option(INFINITRAIN_MACA_WITH_MCCL "Enable MCCL distributed collectives" ON) + +find_library(MACA_RUNTIME_LIBRARY NAMES mcruntime HINTS "${MACA_PATH}/lib" REQUIRED) +find_library(MACA_DNN_LIBRARY NAMES mcdnn HINTS "${MACA_PATH}/lib" REQUIRED) +find_library(MACA_BLAS_LIBRARY NAMES mcblas HINTS "${MACA_PATH}/lib" REQUIRED) + +file(GLOB MACA_DEVICE_SOURCES CONFIGURE_DEPENDS + "${CMAKE_CURRENT_SOURCE_DIR}/src/kernels/*.maca") +set_source_files_properties(${MACA_DEVICE_SOURCES} PROPERTIES + LANGUAGE CXX + COMPILE_OPTIONS "-x;maca" +) + +set(MACA_BACKEND_SOURCES + src/backend.cc + src/runtime/maca_guard_impl.cc + src/runtime/maca_runtime_common.cc + src/kernels/common/gemm.cc + src/kernels/register_maca_kernels.cc + ${MACA_DEVICE_SOURCES} +) + +if(INFINITRAIN_MACA_WITH_MCCL) + find_library(MACA_COMM_LIBRARY NAMES mccl HINTS "${MACA_PATH}/lib" REQUIRED) + list(APPEND MACA_BACKEND_SOURCES + src/ccl/mccl_common.cc + src/ccl/mccl_impl.cc + ) +endif() + +add_library(infini_train_backend_maca STATIC ${MACA_BACKEND_SOURCES}) +add_library(InfiniTrain::Backend::MACA ALIAS infini_train_backend_maca) + +target_include_directories(infini_train_backend_maca + PUBLIC + "$" + PRIVATE + "${CMAKE_CURRENT_SOURCE_DIR}/src" + "${INFINITRAIN_SOURCE_DIR}" + "${MACA_PATH}/include" +) + +target_compile_definitions(infini_train_backend_maca PRIVATE + $<$:USE_MCCL=1> +) + +target_link_libraries(infini_train_backend_maca PUBLIC + InfiniTrain::infini_train + ${MACA_RUNTIME_LIBRARY} + ${MACA_DNN_LIBRARY} + ${MACA_BLAS_LIBRARY} +) +if(INFINITRAIN_MACA_WITH_MCCL) + target_link_libraries(infini_train_backend_maca PUBLIC ${MACA_COMM_LIBRARY}) +endif() + +# Static registries in InfiniTrain and its CPU kernel library are discovered by +# global initialization. Executables should link this target so those objects +# are retained even when no ordinary symbol references them directly. +add_library(infini_train_backend_maca_executable INTERFACE) +add_library(InfiniTrain::Backend::MACAExecutable ALIAS + infini_train_backend_maca_executable) +# Keep target names in the link interface so their usage requirements and build +# dependencies remain attached. This matches InfiniTrain's validated static +# executable link order for MACA. +target_link_libraries(infini_train_backend_maca_executable INTERFACE + "-Wl,--start-group" + "-Wl,--whole-archive" + infini_train + infini_train_cpu_kernels + infini_train_backend_maca + "-Wl,--no-whole-archive" + "-Wl,--end-group" +) + +if(BUILD_TEST) + if(NOT COMMAND infini_train_add_privateuse1_test_suites) + message(FATAL_ERROR + "The selected InfiniTrain revision does not provide PrivateUse1 test suites. " + "Update INFINITRAIN_SOURCE_DIR or the pinned submodule revision.") + endif() + + # The target suffix and CTest label identify this provider; execution still + # uses the backend-neutral PrivateUse1 device and GTest prefix. + infini_train_add_privateuse1_test_suites( + BACKEND_NAME ${INFINITRAIN_BACKEND} + DEVICE_INDEX 0 + LINK_LIBRARIES InfiniTrain::Backend::MACAExecutable + BACKEND_HEADER "infini_train_maca/backend.h" + BACKEND_REGISTRAR infini_train::maca::RegisterBackend + TEST_TIMEOUT 300 + RUN_SERIAL + ) +endif() + +if(INFINITRAIN_BACKENDS_BUILD_EXAMPLES) + add_subdirectory(examples) +endif() diff --git a/backends/maca/cmake/pre_project.cmake b/backends/maca/cmake/pre_project.cmake new file mode 100644 index 0000000..5e759a2 --- /dev/null +++ b/backends/maca/cmake/pre_project.cmake @@ -0,0 +1,82 @@ +set(MACA_PATH "$ENV{MACA_PATH}" CACHE PATH "Path to the MACA SDK") +if(NOT MACA_PATH) + message(FATAL_ERROR + "MACA_PATH is not set. Export MACA_PATH or pass -DMACA_PATH=.") +endif() + +set(_MACA_COMPILER "${MACA_PATH}/mxgpu_llvm/bin/mxcc") +if(NOT EXISTS "${_MACA_COMPILER}") + message(FATAL_ERROR "The MACA compiler was not found at ${_MACA_COMPILER}.") +endif() + +# A compiler can only be selected before the first project() call. Respect an +# explicit toolchain/compiler supplied by an embedding project; otherwise use +# the compiler shipped with the selected MACA SDK. +if(NOT CMAKE_C_COMPILER) + set(CMAKE_C_COMPILER "${_MACA_COMPILER}" CACHE FILEPATH "MACA C compiler") +endif() +if(NOT CMAKE_CXX_COMPILER) + set(CMAKE_CXX_COMPILER "${_MACA_COMPILER}" CACHE FILEPATH "MACA C++ compiler") +endif() +set(CMAKE_CXX_SCAN_FOR_MODULES OFF) + +# FindOpenMP cannot reliably infer mxcc's OpenMP runtime. Seed the variables it +# uses to construct OpenMP::OpenMP_CXX so MACA builds compile with OpenMP and +# link the SDK's ABI-compatible mxomp runtime instead of the system libgomp. +find_library(_MACA_OPENMP_LIBRARY + NAMES mxomp omp iomp5 + HINTS + "${MACA_PATH}/mxgpu_llvm/lib" + "${MACA_PATH}/mxgpu_llvm/lib64" + "${MACA_PATH}/lib" + NO_DEFAULT_PATH + NO_CACHE +) +if(NOT _MACA_OPENMP_LIBRARY) + message(FATAL_ERROR "The MACA OpenMP runtime was not found under ${MACA_PATH}.") +endif() +set(OpenMP_CXX_FLAGS "-fopenmp") +set(OpenMP_CXX_LIB_NAMES "mxomp") +set(OpenMP_mxomp_LIBRARY "${_MACA_OPENMP_LIBRARY}") + +# mxcc cannot reliably run the feature probes used by glog and FindThreads. +# Keep these as directory variables: the InfiniTrain subtree inherits them, +# while an embedding project's cache and sibling directories remain untouched. +# Force FindThreads to select libpthread; claiming libc support clears its link interface. +set(CMAKE_HAVE_LIBC_PTHREAD OFF) +set(CMAKE_HAVE_PTHREADS_CREATE OFF) +set(CMAKE_HAVE_PTHREAD_CREATE ON) +set(HAVE_SYS_TYPES_H 1) +set(HAVE_UNISTD_H 1) +set(HAVE_DLFCN_H 1) +set(HAVE_GLOB_H 1) +set(HAVE_PWD_H 1) +set(HAVE_SYS_TIME_H 1) +set(HAVE_SYS_UTSNAME_H 1) +set(HAVE_SYS_WAIT_H 1) +set(HAVE_SYS_SYSCALL_H 1) +set(HAVE_SYSLOG_H 1) +set(HAVE_UCONTEXT_H 1) +set(HAVE_MODE_T 4) +set(HAVE_HAVE_MODE_T TRUE) +set(HAVE_SSIZE_T 8) +set(HAVE_HAVE_SSIZE_T TRUE) +set(HAVE_PREAD 1) +set(HAVE_PWRITE 1) +set(HAVE_POSIX_FADVISE 1) +set(HAVE_SIGACTION 1) +set(HAVE_SIGALTSTACK 1) +set(HAVE_FCNTL 1) +set(HAVE_DLADDR 1) +set(HAVE___CXA_DEMANGLE 1) + +# Configure only the InfiniTrain subtree selected by this provider. Normal +# variables are sufficient for option() with modern CMake and do not override +# an embedding project's global cache entries. +set(USE_CUDA OFF) +set(USE_NCCL OFF) +set(USE_OMP ON) +set(BUILD_SHARED_LIBS OFF) + +unset(_MACA_COMPILER) +unset(_MACA_OPENMP_LIBRARY) diff --git a/backends/maca/examples/CMakeLists.txt b/backends/maca/examples/CMakeLists.txt new file mode 100644 index 0000000..70c9033 --- /dev/null +++ b/backends/maca/examples/CMakeLists.txt @@ -0,0 +1,36 @@ +function(add_maca_example target_name output_name) + add_executable(${target_name} ${ARGN}) + # Reuse upstream example sources while injecting the selected external + # provider's declaration and registration callback. + target_compile_definitions(${target_name} PRIVATE + INFINITRAIN_EXAMPLE_EXTERNAL_BACKEND_HEADER=\"infini_train_maca/backend.h\" + INFINITRAIN_EXAMPLE_EXTERNAL_BACKEND_REGISTRAR=infini_train::maca::RegisterBackend + ) + target_link_libraries(${target_name} PRIVATE + InfiniTrain::Backend::MACAExecutable) + set_target_properties(${target_name} PROPERTIES + OUTPUT_NAME ${output_name} + RUNTIME_OUTPUT_DIRECTORY "${CMAKE_BINARY_DIR}" + ) +endfunction() + +function(add_maca_model_example model_name) + set(model_dir "${INFINITRAIN_SOURCE_DIR}/example/${model_name}") + set(model_sources + "${INFINITRAIN_SOURCE_DIR}/example/common/tiny_shakespeare_dataset.cc" + "${INFINITRAIN_SOURCE_DIR}/example/common/utils.cc" + "${model_dir}/checkpoint_loader.cc" + ) + if(NOT "${model_name}" STREQUAL "mixtral") + list(APPEND model_sources + "${INFINITRAIN_SOURCE_DIR}/example/common/tokenizer.cc") + endif() + + list(APPEND model_sources "${model_dir}/main.cc") + + add_maca_example("maca_${model_name}" "${model_name}" ${model_sources}) +endfunction() + +add_maca_model_example(gpt2) +add_maca_model_example(llama3) +add_maca_model_example(mixtral) diff --git a/backends/maca/include/infini_train_maca/backend.h b/backends/maca/include/infini_train_maca/backend.h new file mode 100644 index 0000000..6a3a542 --- /dev/null +++ b/backends/maca/include/infini_train_maca/backend.h @@ -0,0 +1,9 @@ +#pragma once + +namespace infini_train::maca { + +// Registers MACA as the process-wide PrivateUse1 provider. Device runtime +// initialization remains lazy until the first device operation. +void RegisterBackend(); + +} // namespace infini_train::maca diff --git a/backends/maca/scripts/run_models_and_profile.bash b/backends/maca/scripts/run_models_and_profile.bash new file mode 100755 index 0000000..101b723 --- /dev/null +++ b/backends/maca/scripts/run_models_and_profile.bash @@ -0,0 +1,21 @@ +#!/usr/bin/env bash + +set -euo pipefail + +SCRIPT_DIR="$(cd -- "$(dirname -- "${BASH_SOURCE[0]}")" && pwd)" +REPOSITORY_ROOT="$(cd -- "${SCRIPT_DIR}/../../.." && pwd)" +INFINITRAIN_SOURCE_DIR="${INFINITRAIN_SOURCE_DIR:-${REPOSITORY_ROOT}/third_party/InfiniTrain}" +UPSTREAM_RUNNER="${INFINITRAIN_SOURCE_DIR}/scripts/run_models_and_profile.bash" + +if [[ ! -x "${UPSTREAM_RUNNER}" ]]; then + echo "Error: InfiniTrain test runner was not found at ${UPSTREAM_RUNNER}." >&2 + echo "Initialize the submodule or set INFINITRAIN_SOURCE_DIR." >&2 + exit 1 +fi + +# InfiniTrain's runner resolves build and log paths from the working directory. +# Anchor those paths at this repository and supply MACA's provider configuration. +cd "${REPOSITORY_ROOT}" +exec "${UPSTREAM_RUNNER}" \ + --test-config "${SCRIPT_DIR}/test_config_maca.json" \ + "$@" diff --git a/backends/maca/scripts/test_config_maca.json b/backends/maca/scripts/test_config_maca.json new file mode 100644 index 0000000..6877ea4 --- /dev/null +++ b/backends/maca/scripts/test_config_maca.json @@ -0,0 +1,693 @@ +{ + "variables": { + "BUILD_DIR": "build/maca", + "GPT2_INPUT_BIN": "/nfs/InfiniTrain-dev/data/llmc/gpt2/tinyshakespeare/tiny_shakespeare_train.bin", + "GPT2_LLMC_FILEPATH": "/nfs/InfiniTrain-dev/data/llmc/gpt2/gpt2_124M.bin", + "LLAMA3_INPUT_BIN": "/nfs/InfiniTrain-dev/data/llmc/llama3/tinyshakespeare/tiny_shakespeare_train.bin", + "LLAMA3_LLMC_FILEPATH": "/nfs/InfiniTrain-dev/data/llmc/llama3/llama3.2_1B_fp32.bin", + "MIXTRAL_INPUT_BIN": "/nfs/InfiniTrain-dev/data/llmc/llama3/tinyshakespeare/tiny_shakespeare_train.bin", + "MIXTRAL_LLMC_FILEPATH": "/nfs/InfiniTrain-dev/data/llmc/mixtral/mixtral_megatron_export.bin", + "PROFILE_LOG_DIR": "./profile_logs", + "LOG_DIR": "./logs", + "COMPARE_LOG_DIR": "", + "RUN_CTEST": "true", + "RUN_PROFILE_TEST": "true", + "CKPT_ROOT_DIR": "/tmp/infini_train_maca_ckpt", + "DEVICE_BACKEND": "maca", + "GPT2_TEST_GROUPS": "basic,zero,lora,8_proc", + "LLAMA3_TEST_GROUPS": "basic,zero,lora,8_proc", + "MIXTRAL_TEST_GROUPS": "moe" + }, + "basic_compile_commands": [ + { + "id": "build_maca", + "cmd": "cmake -DINFINITRAIN_BACKEND=maca -DINFINITRAIN_MACA_WITH_MCCL=ON ../.. && cmake --build . -j" + } + ], + "test_groups": [ + { + "tag": "basic", + "tests": [ + { + "id": "1", + "args": { + "dtype": "float32" + } + }, + { + "id": "1_bfloat16", + "args": { + "dtype": "bfloat16" + } + }, + { + "id": "2", + "args": { + "dtype": "float32", + "num_iteration": 10, + "batch_size": 80, + "total_batch_size": 5120 + } + }, + { + "id": "2_bfloat16", + "args": { + "dtype": "bfloat16", + "num_iteration": 10, + "batch_size": 80, + "total_batch_size": 5120 + } + }, + { + "id": "3", + "args": { + "dtype": "float32", + "nthread_per_process": 8, + "num_iteration": 10, + "batch_size": 10, + "total_batch_size": 5120 + } + }, + { + "id": "3_bfloat16", + "args": { + "dtype": "bfloat16", + "nthread_per_process": 8, + "num_iteration": 10, + "batch_size": 10, + "total_batch_size": 5120 + } + }, + { + "id": "4", + "args": { + "dtype": "float32", + "nthread_per_process": 8, + "num_iteration": 10, + "batch_size": 40, + "total_batch_size": 5120, + "tensor_parallel": 4 + } + }, + { + "id": "4_bfloat16", + "args": { + "dtype": "bfloat16", + "nthread_per_process": 8, + "num_iteration": 10, + "batch_size": 40, + "total_batch_size": 5120, + "tensor_parallel": 4 + } + }, + { + "id": "5", + "args": { + "dtype": "float32", + "nthread_per_process": 8, + "num_iteration": 10, + "batch_size": 40, + "total_batch_size": 5120, + "tensor_parallel": 4, + "sequence_parallel": true + } + }, + { + "id": "5_bfloat16", + "args": { + "dtype": "bfloat16", + "nthread_per_process": 8, + "num_iteration": 10, + "batch_size": 40, + "total_batch_size": 5120, + "tensor_parallel": 4, + "sequence_parallel": true + } + }, + { + "id": "6", + "args": { + "dtype": "float32", + "nthread_per_process": 8, + "num_iteration": 10, + "batch_size": 10, + "total_batch_size": 5120, + "pipeline_parallel": 8 + } + }, + { + "id": "6_bfloat16", + "args": { + "dtype": "bfloat16", + "nthread_per_process": 8, + "num_iteration": 10, + "batch_size": 10, + "total_batch_size": 5120, + "pipeline_parallel": 8 + } + }, + { + "id": "7", + "args": { + "dtype": "float32", + "nthread_per_process": 4, + "num_iteration": 10, + "batch_size": 10, + "total_batch_size": 5120, + "pipeline_parallel": 4, + "virtual_pipeline_parallel": 2 + } + }, + { + "id": "7_bfloat16", + "args": { + "dtype": "bfloat16", + "nthread_per_process": 4, + "num_iteration": 10, + "batch_size": 10, + "total_batch_size": 5120, + "pipeline_parallel": 4, + "virtual_pipeline_parallel": 2 + } + }, + { + "id": "8", + "args": { + "dtype": "float32", + "nthread_per_process": 8, + "num_iteration": 10, + "batch_size": 40, + "total_batch_size": 5120, + "tensor_parallel": 2, + "sequence_parallel": true, + "pipeline_parallel": 2, + "virtual_pipeline_parallel": 2 + } + }, + { + "id": "8_bfloat16", + "args": { + "dtype": "bfloat16", + "nthread_per_process": 8, + "num_iteration": 10, + "batch_size": 40, + "total_batch_size": 5120, + "tensor_parallel": 2, + "sequence_parallel": true, + "pipeline_parallel": 2, + "virtual_pipeline_parallel": 2 + } + } + ] + }, + { + "tag": "zero", + "tests": [ + { + "id": "3_distopt", + "args": { + "dtype": "float32", + "nthread_per_process": 8, + "num_iteration": 10, + "batch_size": 10, + "total_batch_size": 5120, + "zero_stage": 1 + } + }, + { + "id": "3_bfloat16_distopt", + "args": { + "dtype": "bfloat16", + "nthread_per_process": 8, + "num_iteration": 10, + "batch_size": 10, + "total_batch_size": 5120, + "zero_stage": 1 + } + }, + { + "id": "4_distopt", + "args": { + "dtype": "float32", + "nthread_per_process": 8, + "num_iteration": 10, + "batch_size": 40, + "total_batch_size": 5120, + "tensor_parallel": 4, + "zero_stage": 1 + } + }, + { + "id": "4_bfloat16_distopt", + "args": { + "dtype": "bfloat16", + "nthread_per_process": 8, + "num_iteration": 10, + "batch_size": 40, + "total_batch_size": 5120, + "tensor_parallel": 4, + "zero_stage": 1 + } + }, + { + "id": "5_distopt", + "args": { + "dtype": "float32", + "nthread_per_process": 8, + "num_iteration": 10, + "batch_size": 40, + "total_batch_size": 5120, + "tensor_parallel": 4, + "sequence_parallel": true, + "zero_stage": 1 + } + }, + { + "id": "5_bfloat16_distopt", + "args": { + "dtype": "bfloat16", + "nthread_per_process": 8, + "num_iteration": 10, + "batch_size": 40, + "total_batch_size": 5120, + "tensor_parallel": 4, + "sequence_parallel": true, + "zero_stage": 1 + } + }, + { + "id": "8_distopt", + "args": { + "dtype": "float32", + "nthread_per_process": 8, + "num_iteration": 10, + "batch_size": 40, + "total_batch_size": 5120, + "tensor_parallel": 2, + "sequence_parallel": true, + "pipeline_parallel": 2, + "virtual_pipeline_parallel": 2, + "zero_stage": 1 + } + }, + { + "id": "8_bfloat16_distopt", + "args": { + "dtype": "bfloat16", + "nthread_per_process": 8, + "num_iteration": 10, + "batch_size": 40, + "total_batch_size": 5120, + "tensor_parallel": 2, + "sequence_parallel": true, + "pipeline_parallel": 2, + "virtual_pipeline_parallel": 2, + "zero_stage": 1 + } + } + ] + }, + { + "tag": "lora", + "tests": [ + { + "id": "1_lora", + "args": { + "dtype": "float32", + "lora_rank": 8, + "lora_alpha": 16.0, + "lora_target_modules": "c_attn,attn.c_proj" + } + }, + { + "id": "1_lora_bfloat16", + "args": { + "dtype": "bfloat16", + "lora_rank": 8, + "lora_alpha": 16.0, + "lora_target_modules": "c_attn,attn.c_proj" + } + }, + { + "id": "2_lora", + "args": { + "dtype": "float32", + "num_iteration": 10, + "batch_size": 80, + "total_batch_size": 5120, + "lora_rank": 8, + "lora_alpha": 16.0, + "lora_target_modules": "c_attn,attn.c_proj" + } + }, + { + "id": "2_lora_bfloat16", + "args": { + "dtype": "bfloat16", + "num_iteration": 10, + "batch_size": 80, + "total_batch_size": 5120, + "lora_rank": 8, + "lora_alpha": 16.0, + "lora_target_modules": "c_attn,attn.c_proj" + } + }, + { + "id": "3_lora", + "args": { + "dtype": "float32", + "nthread_per_process": 8, + "num_iteration": 10, + "batch_size": 10, + "total_batch_size": 5120, + "lora_rank": 8, + "lora_alpha": 16.0, + "lora_target_modules": "c_attn,attn.c_proj" + } + }, + { + "id": "3_lora_bfloat16", + "args": { + "dtype": "bfloat16", + "nthread_per_process": 8, + "num_iteration": 10, + "batch_size": 10, + "total_batch_size": 5120, + "lora_rank": 8, + "lora_alpha": 16.0, + "lora_target_modules": "c_attn,attn.c_proj" + } + }, + { + "id": "4_lora", + "args": { + "dtype": "float32", + "nthread_per_process": 8, + "num_iteration": 10, + "batch_size": 40, + "total_batch_size": 5120, + "tensor_parallel": 4, + "lora_rank": 4, + "lora_alpha": 8.0, + "lora_target_modules": "c_attn,c_fc,c_proj" + } + }, + { + "id": "4_lora_bfloat16", + "args": { + "dtype": "bfloat16", + "nthread_per_process": 8, + "num_iteration": 10, + "batch_size": 40, + "total_batch_size": 5120, + "tensor_parallel": 4, + "lora_rank": 16, + "lora_alpha": 32.0, + "lora_target_modules": "c_attn,c_fc,c_proj" + } + }, + { + "id": "5_lora", + "args": { + "dtype": "float32", + "nthread_per_process": 8, + "num_iteration": 10, + "batch_size": 40, + "total_batch_size": 5120, + "tensor_parallel": 4, + "sequence_parallel": true, + "lora_rank": 4, + "lora_alpha": 8.0, + "lora_target_modules": "attn.c_proj,c_fc,c_proj" + } + }, + { + "id": "5_lora_bfloat16", + "args": { + "dtype": "bfloat16", + "nthread_per_process": 8, + "num_iteration": 10, + "batch_size": 40, + "total_batch_size": 5120, + "tensor_parallel": 4, + "sequence_parallel": true, + "lora_rank": 16, + "lora_alpha": 32.0, + "lora_target_modules": "attn.c_proj,c_fc,c_proj" + } + }, + { + "id": "6_lora", + "args": { + "dtype": "float32", + "nthread_per_process": 8, + "num_iteration": 10, + "batch_size": 10, + "total_batch_size": 5120, + "pipeline_parallel": 8, + "lora_rank": 4, + "lora_alpha": 8.0, + "lora_target_modules": "c_attn,attn.c_proj,c_fc" + } + }, + { + "id": "6_lora_bfloat16", + "args": { + "dtype": "bfloat16", + "nthread_per_process": 8, + "num_iteration": 10, + "batch_size": 10, + "total_batch_size": 5120, + "pipeline_parallel": 8, + "lora_rank": 16, + "lora_alpha": 32.0, + "lora_target_modules": "c_attn,attn.c_proj,c_fc" + } + }, + { + "id": "7_lora", + "args": { + "dtype": "float32", + "nthread_per_process": 4, + "num_iteration": 10, + "batch_size": 10, + "total_batch_size": 5120, + "pipeline_parallel": 4, + "virtual_pipeline_parallel": 2, + "lora_rank": 4, + "lora_alpha": 8.0, + "lora_target_modules": "c_attn,c_proj" + } + }, + { + "id": "7_lora_bfloat16", + "args": { + "dtype": "bfloat16", + "nthread_per_process": 4, + "num_iteration": 10, + "batch_size": 10, + "total_batch_size": 5120, + "pipeline_parallel": 4, + "virtual_pipeline_parallel": 2, + "lora_rank": 16, + "lora_alpha": 32.0, + "lora_target_modules": "c_attn,c_proj" + } + }, + { + "id": "8_lora", + "args": { + "dtype": "float32", + "nthread_per_process": 8, + "num_iteration": 10, + "batch_size": 40, + "total_batch_size": 5120, + "tensor_parallel": 2, + "sequence_parallel": true, + "pipeline_parallel": 2, + "virtual_pipeline_parallel": 2, + "lora_rank": 4, + "lora_alpha": 8.0, + "lora_target_modules": "c_fc,c_proj" + } + }, + { + "id": "8_lora_bfloat16", + "args": { + "dtype": "bfloat16", + "nthread_per_process": 8, + "num_iteration": 10, + "batch_size": 40, + "total_batch_size": 5120, + "tensor_parallel": 2, + "sequence_parallel": true, + "pipeline_parallel": 2, + "virtual_pipeline_parallel": 2, + "lora_rank": 16, + "lora_alpha": 32.0, + "lora_target_modules": "c_fc,c_proj" + } + } + ] + }, + { + "tag": "moe", + "tests": [ + { + "id": "1", + "args": { + "dtype": "float32" + } + }, + { + "id": "1_bfloat16", + "args": { + "dtype": "bfloat16" + } + }, + { + "id": "2", + "args": { + "dtype": "float32", + "num_iteration": 10, + "global_batch_size": 80, + "micro_batch_size": 80 + } + }, + { + "id": "2_bfloat16", + "args": { + "dtype": "bfloat16", + "num_iteration": 10, + "global_batch_size": 80, + "micro_batch_size": 80 + } + } + ] + }, + { + "tag": "8_proc", + "infini_run_args": { + "nproc_per_node": 8 + }, + "tests": [ + { + "id": "3_8proc", + "args": { + "dtype": "float32", + "nthread_per_process": 1, + "num_iteration": 10, + "batch_size": 10, + "total_batch_size": 5120 + } + }, + { + "id": "3_bfloat16_8proc", + "args": { + "dtype": "bfloat16", + "nthread_per_process": 1, + "num_iteration": 10, + "batch_size": 10, + "total_batch_size": 5120 + } + }, + { + "id": "4_8proc", + "args": { + "dtype": "float32", + "nthread_per_process": 1, + "num_iteration": 10, + "batch_size": 40, + "total_batch_size": 5120, + "tensor_parallel": 4 + } + }, + { + "id": "4_bfloat16_8proc", + "args": { + "dtype": "bfloat16", + "nthread_per_process": 1, + "num_iteration": 10, + "batch_size": 40, + "total_batch_size": 5120, + "tensor_parallel": 4 + } + }, + { + "id": "5_8proc", + "args": { + "dtype": "float32", + "nthread_per_process": 1, + "num_iteration": 10, + "batch_size": 40, + "total_batch_size": 5120, + "tensor_parallel": 4, + "sequence_parallel": true + } + }, + { + "id": "5_bfloat16_8proc", + "args": { + "dtype": "bfloat16", + "nthread_per_process": 1, + "num_iteration": 10, + "batch_size": 40, + "total_batch_size": 5120, + "tensor_parallel": 4, + "sequence_parallel": true + } + }, + { + "id": "6_8proc", + "args": { + "dtype": "float32", + "nthread_per_process": 1, + "num_iteration": 10, + "batch_size": 10, + "total_batch_size": 5120, + "pipeline_parallel": 8 + } + }, + { + "id": "6_bfloat16_8proc", + "args": { + "dtype": "bfloat16", + "nthread_per_process": 1, + "num_iteration": 10, + "batch_size": 10, + "total_batch_size": 5120, + "pipeline_parallel": 8 + } + }, + { + "id": "8_8proc", + "args": { + "dtype": "float32", + "nthread_per_process": 1, + "num_iteration": 10, + "batch_size": 40, + "total_batch_size": 5120, + "tensor_parallel": 2, + "sequence_parallel": true, + "pipeline_parallel": 2, + "virtual_pipeline_parallel": 2 + } + }, + { + "id": "8_bfloat16_8proc", + "args": { + "dtype": "bfloat16", + "nthread_per_process": 1, + "num_iteration": 10, + "batch_size": 40, + "total_batch_size": 5120, + "tensor_parallel": 2, + "sequence_parallel": true, + "pipeline_parallel": 2, + "virtual_pipeline_parallel": 2 + } + } + ] + } + ] +} diff --git a/backends/maca/src/backend.cc b/backends/maca/src/backend.cc new file mode 100644 index 0000000..a99e298 --- /dev/null +++ b/backends/maca/src/backend.cc @@ -0,0 +1,30 @@ +#include "infini_train_maca/backend.h" + +#include + +#include "infini_train/include/core/privateuse1_backend.h" + +#include "kernels/register_maca_kernels.h" +#include "runtime/maca_guard_impl.h" +#ifdef USE_MCCL +#include "ccl/mccl_impl.h" +#endif + +namespace infini_train::maca { + +void RegisterBackend() { + static std::once_flag once; + std::call_once(once, []() { + core::PrivateUse1BackendRegistration registration; + registration.name = "maca"; + registration.default_autocast_dtype = DataType::kBFLOAT16; + registration.register_runtime = &core::maca::RegisterMacaRuntime; + registration.register_kernels = &kernels::maca::RegisterMacaKernels; +#ifdef USE_MCCL + registration.register_ccl = &core::maca::RegisterMcclBackend; +#endif + core::RegisterPrivateUse1Backend(registration); + }); +} + +} // namespace infini_train::maca diff --git a/backends/maca/src/ccl/mccl_common.cc b/backends/maca/src/ccl/mccl_common.cc new file mode 100644 index 0000000..dccd0c9 --- /dev/null +++ b/backends/maca/src/ccl/mccl_common.cc @@ -0,0 +1,35 @@ +#include "ccl/mccl_common.h" + +#include + +#include "glog/logging.h" + +namespace infini_train::core::maca { + +McclComm::McclComm() = default; + +McclComm::McclComm(mcclComm_t comm) : mccl_comm_(comm) {} + +mcclComm_t McclComm::mccl_comm() const { return mccl_comm_; } + +void McclComm::set_mccl_comm(mcclComm_t comm) { mccl_comm_ = comm; } + +McclUniqueId::McclUniqueId() = default; + +McclUniqueId::McclUniqueId(const mcclUniqueId &id) : id_(id) {} + +size_t McclUniqueId::Size() const { return sizeof(id_); } + +const void *McclUniqueId::Data() const { return &id_; } + +void McclUniqueId::Load(const void *src, size_t size) { + CHECK_NOTNULL(src); + CHECK_EQ(size, sizeof(id_)); + std::memcpy(&id_, src, sizeof(id_)); +} + +mcclUniqueId *McclUniqueId::mccl_unique_id() { return &id_; } + +const mcclUniqueId *McclUniqueId::mccl_unique_id() const { return &id_; } + +} // namespace infini_train::core::maca diff --git a/backends/maca/src/ccl/mccl_common.h b/backends/maca/src/ccl/mccl_common.h new file mode 100644 index 0000000..d30596d --- /dev/null +++ b/backends/maca/src/ccl/mccl_common.h @@ -0,0 +1,39 @@ +#pragma once + +#include + +#include + +#include "infini_train/include/core/ccl/ccl_common.h" + +namespace infini_train::core::maca { + +class McclComm final : public CclComm { +public: + McclComm(); + explicit McclComm(mcclComm_t comm); + + mcclComm_t mccl_comm() const; + void set_mccl_comm(mcclComm_t comm); + +private: + mcclComm_t mccl_comm_ = nullptr; +}; + +class McclUniqueId final : public CclUniqueId { +public: + McclUniqueId(); + explicit McclUniqueId(const mcclUniqueId &id); + + size_t Size() const override; + const void *Data() const override; + void Load(const void *src, size_t size) override; + + mcclUniqueId *mccl_unique_id(); + const mcclUniqueId *mccl_unique_id() const; + +private: + mcclUniqueId id_; +}; + +} // namespace infini_train::core::maca diff --git a/backends/maca/src/ccl/mccl_impl.cc b/backends/maca/src/ccl/mccl_impl.cc new file mode 100644 index 0000000..a355a41 --- /dev/null +++ b/backends/maca/src/ccl/mccl_impl.cc @@ -0,0 +1,162 @@ +#include "ccl/mccl_impl.h" + +#include +#include +#include +#include + +#include "glog/logging.h" + +#include "infini_train/include/core/runtime/runtime_common.h" +#include "infini_train/include/device.h" + +#include "ccl/mccl_common.h" +#include "common/common_maca.h" +#include "runtime/maca_runtime_common.h" + +namespace infini_train::core::maca { +namespace { + +inline const std::unordered_map kMcclDtypeMap = { + {DataType::kUINT8, mcclUint8}, {DataType::kINT8, mcclInt8}, {DataType::kUINT32, mcclUint32}, + {DataType::kINT32, mcclInt32}, {DataType::kUINT64, mcclUint64}, {DataType::kINT64, mcclInt64}, + {DataType::kBFLOAT16, mcclBfloat16}, {DataType::kFLOAT16, mcclHalf}, {DataType::kFLOAT32, mcclFloat32}, + {DataType::kFLOAT64, mcclFloat64}, +}; + +inline const std::unordered_map kMcclReduceOpMap = { + {nn::parallel::function::ReduceOpType::kSum, mcclSum}, {nn::parallel::function::ReduceOpType::kProd, mcclProd}, + {nn::parallel::function::ReduceOpType::kMin, mcclMin}, {nn::parallel::function::ReduceOpType::kMax, mcclMax}, + {nn::parallel::function::ReduceOpType::kAvg, mcclAvg}, +}; + +inline mcclComm_t GetMcclComm(const CclComm *comm) { + auto *mccl_comm = dynamic_cast(comm); + CHECK_NOTNULL(mccl_comm); + return mccl_comm->mccl_comm(); +} + +inline void SetMcclComm(CclComm *comm, mcclComm_t mccl_comm) { + auto *typed_comm = dynamic_cast(comm); + CHECK_NOTNULL(typed_comm); + typed_comm->set_mccl_comm(mccl_comm); +} + +inline const mcclUniqueId &GetMcclUniqueId(const CclUniqueId &unique_id) { + auto *mccl_unique_id = dynamic_cast(&unique_id); + CHECK_NOTNULL(mccl_unique_id); + return *mccl_unique_id->mccl_unique_id(); +} + +inline mcStream_t GetMacaStream(Stream *stream) { + auto *maca_stream = dynamic_cast(stream); + CHECK_NOTNULL(maca_stream); + return maca_stream->maca_stream(); +} + +} // namespace + +Device::DeviceType McclImpl::Type() const { return Device::DeviceType::kPrivateUse1; } + +void McclImpl::GroupStart() const { MCCL_CHECK(mcclGroupStart()); } + +void McclImpl::GroupEnd() const { MCCL_CHECK(mcclGroupEnd()); } + +void McclImpl::GetAsyncError(const CclComm *comm, CclStatus *async_error) const { + mcclResult_t mccl_async_error = mcclSuccess; + MCCL_CHECK(mcclCommGetAsyncError(GetMcclComm(comm), &mccl_async_error)); + if (async_error != nullptr) { + *async_error = (mccl_async_error == mcclSuccess) ? CclStatus::kSuccess : CclStatus::kError; + } +} + +void McclImpl::GetUniqueId(CclUniqueId **unique_id) const { + CHECK_NOTNULL(unique_id); + if (*unique_id == nullptr) { + *unique_id = new McclUniqueId(); + } + auto *mccl_unique_id = dynamic_cast(*unique_id); + CHECK_NOTNULL(mccl_unique_id); + MCCL_CHECK(mcclGetUniqueId(mccl_unique_id->mccl_unique_id())); +} + +void McclImpl::CommInitAll(CclComm **comms, int ndev, const int *devlist) const { + CHECK_NOTNULL(comms); + CHECK_GT(ndev, 0); + CHECK_NOTNULL(devlist); + + std::vector mccl_comms(static_cast(ndev), nullptr); + MCCL_CHECK(mcclCommInitAll(mccl_comms.data(), ndev, devlist)); + for (int i = 0; i < ndev; ++i) { + if (comms[i] == nullptr) { + comms[i] = new McclComm(); + } + SetMcclComm(comms[i], mccl_comms[static_cast(i)]); + } +} + +void McclImpl::CommInitRank(CclComm **comm, int nranks, const CclUniqueId &unique_id, int rank) const { + CHECK_NOTNULL(comm); + CHECK_GT(nranks, 0); + + if (*comm == nullptr) { + *comm = new McclComm(); + } + + mcclComm_t mccl_comm = nullptr; + MCCL_CHECK(mcclCommInitRank(&mccl_comm, nranks, GetMcclUniqueId(unique_id), rank)); + SetMcclComm(*comm, mccl_comm); +} + +void McclImpl::CommDestroy(CclComm *comm) const { + if (comm == nullptr) { + return; + } + MCCL_CHECK(mcclCommDestroy(GetMcclComm(comm))); + SetMcclComm(comm, nullptr); +} + +void McclImpl::AllReduce(const void *sendbuff, void *recvbuff, size_t count, DataType dtype, + nn::parallel::function::ReduceOpType reduce_op, const CclComm *comm, Stream *stream) const { + MCCL_CHECK(mcclAllReduce(sendbuff, recvbuff, count, kMcclDtypeMap.at(dtype), kMcclReduceOpMap.at(reduce_op), + GetMcclComm(comm), GetMacaStream(stream))); +} + +void McclImpl::Broadcast(const void *sendbuff, void *recvbuff, size_t count, DataType dtype, int root, + const CclComm *comm, Stream *stream) const { + MCCL_CHECK(mcclBroadcast(sendbuff, recvbuff, count, kMcclDtypeMap.at(dtype), root, GetMcclComm(comm), + GetMacaStream(stream))); +} + +void McclImpl::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 { + MCCL_CHECK(mcclReduce(sendbuff, recvbuff, count, kMcclDtypeMap.at(dtype), kMcclReduceOpMap.at(reduce_op), root, + GetMcclComm(comm), GetMacaStream(stream))); +} + +void McclImpl::AllGather(const void *sendbuff, void *recvbuff, size_t count, DataType dtype, const CclComm *comm, + Stream *stream) const { + MCCL_CHECK( + mcclAllGather(sendbuff, recvbuff, count, kMcclDtypeMap.at(dtype), GetMcclComm(comm), GetMacaStream(stream))); +} + +void McclImpl::ReduceScatter(const void *sendbuff, void *recvbuff, size_t recv_count, DataType dtype, + nn::parallel::function::ReduceOpType reduce_op, const CclComm *comm, + Stream *stream) const { + MCCL_CHECK(mcclReduceScatter(sendbuff, recvbuff, recv_count, kMcclDtypeMap.at(dtype), + kMcclReduceOpMap.at(reduce_op), GetMcclComm(comm), GetMacaStream(stream))); +} + +void McclImpl::Send(const void *buff, size_t count, DataType dtype, int peer, const CclComm *comm, + Stream *stream) const { + MCCL_CHECK(mcclSend(buff, count, kMcclDtypeMap.at(dtype), peer, GetMcclComm(comm), GetMacaStream(stream))); +} + +void McclImpl::Recv(void *buff, size_t count, DataType dtype, int peer, const CclComm *comm, Stream *stream) const { + MCCL_CHECK(mcclRecv(buff, count, kMcclDtypeMap.at(dtype), peer, GetMcclComm(comm), GetMacaStream(stream))); +} + +void RegisterMcclBackend() { INFINI_TRAIN_REGISTER_CCL_IMPL(Device::DeviceType::kPrivateUse1, McclImpl) } + +} // namespace infini_train::core::maca diff --git a/backends/maca/src/ccl/mccl_impl.h b/backends/maca/src/ccl/mccl_impl.h new file mode 100644 index 0000000..a97488a --- /dev/null +++ b/backends/maca/src/ccl/mccl_impl.h @@ -0,0 +1,52 @@ +#pragma once + +#include + +#include "infini_train/include/core/ccl/ccl.h" + +namespace infini_train::core::maca { + +void RegisterMcclBackend(); + +class McclImpl 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::maca diff --git a/backends/maca/src/common/common_maca.h b/backends/maca/src/common/common_maca.h new file mode 100644 index 0000000..6df6759 --- /dev/null +++ b/backends/maca/src/common/common_maca.h @@ -0,0 +1,44 @@ +#pragma once + +#include +#include +#include + +#ifdef USE_MCCL +#include +#endif + +#include "glog/logging.h" + +namespace infini_train::common::maca { + +// MACA runtime and library call checks used by this backend implementation. +#define MACA_CHECK(call) \ + do { \ + mcError_t _maca_status = (call); \ + if (_maca_status != mcSuccess) { \ + LOG(FATAL) << "MACA Error: " << mcGetErrorString(_maca_status) << " at " << __FILE__ << ":" << __LINE__; \ + } \ + } while (0) + +#define MCBLAS_CHECK(call) \ + do { \ + mcblasStatus_t _mcblas_status = (call); \ + if (_mcblas_status != MCBLAS_STATUS_SUCCESS) { \ + LOG(FATAL) << "MCBLAS Error: " << mcblasGetStatusString(_mcblas_status) << " at " << __FILE__ << ":" \ + << __LINE__; \ + } \ + } while (0) + +#ifdef USE_MCCL +#define MCCL_CHECK(expr) \ + do { \ + mcclResult_t _status = (expr); \ + if (_status != mcclSuccess) { \ + LOG(FATAL) << "MCCL error: " << mcclGetErrorString(_status) << " at " << __FILE__ << ":" << __LINE__ \ + << " (" << #expr << ")"; \ + } \ + } while (0) +#endif + +} // namespace infini_train::common::maca diff --git a/backends/maca/src/common/cub_compat.cuh b/backends/maca/src/common/cub_compat.cuh new file mode 100644 index 0000000..52518ff --- /dev/null +++ b/backends/maca/src/common/cub_compat.cuh @@ -0,0 +1,13 @@ +#pragma once + +#include + +namespace infini_train::kernels::maca { + +// MACA ships a CUB compatible with the pre-2.8 API (cub::Sum/Max/Min). +// Match the CUDA reduction aliases used by the shared kernel conventions. +using CubSumOp = cub::Sum; +using CubMaxOp = cub::Max; +using CubMinOp = cub::Min; + +} // namespace infini_train::kernels::maca diff --git a/backends/maca/src/common/kernel_helper.cuh b/backends/maca/src/common/kernel_helper.cuh new file mode 100644 index 0000000..0c43026 --- /dev/null +++ b/backends/maca/src/common/kernel_helper.cuh @@ -0,0 +1,320 @@ +#pragma once + +#include +#include +#include +#include +#include + +#include +#include + +namespace infini_train::common::maca { +/** + * Converts a value between arbitrary types with specialized handling for + * MACA 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<__half>(3.14f); // float -> half (MACA intrinsic) + * float f = Cast(h); // half -> float (MACA intrinsic) + * int i = Cast(2.718); // double -> int (standard cast) + */ +// TODO(zbl): add support for half and __maca_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>; + + // __maca_bfloat16 conversions + if constexpr (std::is_same_v) { + if constexpr (std::is_same_v) { + return __bfloat162float(x); + } else if constexpr (std::is_same_v) { + return static_cast(__bfloat162float(x)); + } else if constexpr (std::is_same_v) { + return __half(__bfloat162float(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 __maca_bfloat16(__half2float(x)); + } + } + // float conversions to reduced precision + else if constexpr (std::is_same_v) { + if constexpr (std::is_same_v) { + return __float2bfloat16(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 __double2bfloat16(x); + } else if constexpr (std::is_same_v) { + return __double2half(x); + } + } + // Fallback for all other conversions + if constexpr (std::is_same_v || std::is_same_v + || std::is_same_v || std::is_same_v) { + return (DST)(static_cast(std::forward(x))); + ; + } else { + return static_cast(std::forward(x)); + ; + } +} + +template __device__ __forceinline__ T Neg(const T &x) { + if constexpr (std::is_same_v || 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 __hdiv(__float2bfloat16(1.0f), 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 __float2bfloat16(__sinf(__bfloat162float(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 __float2bfloat16(__cosf(__bfloat162float(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 __float2half(tanhf(__half2float(x))); + } else if constexpr (std::is_same_v) { + return __float2bfloat16(tanhf(__bfloat162float(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_ = __bfloat162float(x); + float exponent_ = __bfloat162float(exponent); + float ans_f = __powf(x_, exponent_); + return __float2bfloat16(__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 __float2bfloat16(rsqrtf(__bfloat162float(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 || 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 __float2bfloat16(__logf(__bfloat162float(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 || 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 || 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 || 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 || 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 || 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 || 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 || 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 __float2bfloat16(__fmaf_rn(__bfloat162float(x), __bfloat162float(y), __bfloat162float(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) { + __maca_bfloat16 *target_addr = tensor + index; + bool low_byte = ((reinterpret_cast(target_addr) & (sizeof(__maca_bfloat162) - 1)) == 0); + + if (low_byte && index < (num_elements - 1)) { + __maca_bfloat162 value2 = __halves2bfloat162(value, __maca_bfloat16(0.0f)); + atomicAdd(reinterpret_cast<__maca_bfloat162 *>(target_addr), value2); + + } else if (!low_byte && index > 0) { + __maca_bfloat162 value2 = __halves2bfloat162(__maca_bfloat16(0.0f), value); + atomicAdd(reinterpret_cast<__maca_bfloat162 *>(target_addr - 1), value2); + + } else { + atomicAdd(target_addr, 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::maca diff --git a/backends/maca/src/kernels/accumulate_grad.maca b/backends/maca/src/kernels/accumulate_grad.maca new file mode 100644 index 0000000..7cf2d5a --- /dev/null +++ b/backends/maca/src/kernels/accumulate_grad.maca @@ -0,0 +1,100 @@ +#include +#include +#include +#include + +#include "common/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 "runtime/maca_dispatch.h" +#include "runtime/maca_runtime_common.h" + +namespace infini_train::kernels::maca { + +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::maca::Mul(grad_ptr[idx], common::maca::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 &maca_stream = dynamic_cast( + infini_train::core::GetDeviceGuardImpl(device.type())->GetStream(device)) + ->maca_stream(); + + core::maca::DispatchMacaFunc( + gradient->Dtype(), + [=]() { + AccumulateGradKernel<<>>( + static_cast(gradient->DataPtr()), rate, static_cast(tensor->DataPtr()), num_elements); + }, + "MACA 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::maca::Fma(common::maca::Cast(beta1), m_data[idx], + common::maca::Cast(1 - beta1) * grad_data[idx]); + v_data[idx] = common::maca::Fma(common::maca::Cast(beta2), v_data[idx], + common::maca::Cast(1 - beta2) * grad_data[idx] * grad_data[idx]); + + const float m_hat = common::maca::Cast(m_data[idx]) / bias_correction_m; + const float v_hat = common::maca::Cast(v_data[idx]) / bias_correction_v; + + param_data[idx] = common::maca::Sub( + param_data[idx], common::maca::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 &maca_stream = dynamic_cast( + infini_train::core::GetDeviceGuardImpl(device.type())->GetStream(device)) + ->maca_stream(); + + core::maca::DispatchMacaFunc( + 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); + }, + "MACA AdamAccumulateGrad"); +} +} // namespace infini_train::kernels::maca + +namespace infini_train::kernels::maca { + +void RegisterAccumulateGradKernels() { + REGISTER_KERNEL(infini_train::Device::DeviceType::kPrivateUse1, AccumulateGrad, + infini_train::kernels::maca::AccumulateGrad) + REGISTER_KERNEL(infini_train::Device::DeviceType::kPrivateUse1, AdamAccumulateGrad, + infini_train::kernels::maca::AdamAccumulateGrad) +} + +} // namespace infini_train::kernels::maca diff --git a/backends/maca/src/kernels/cast.maca b/backends/maca/src/kernels/cast.maca new file mode 100644 index 0000000..5a42b33 --- /dev/null +++ b/backends/maca/src/kernels/cast.maca @@ -0,0 +1,60 @@ +#include +#include + +#include "infini_train/include/common/common.h" +#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 "common/kernel_helper.cuh" +#include "runtime/maca_dispatch.h" +#include "runtime/maca_runtime_common.h" + +namespace infini_train::kernels::maca { + +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::maca::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 &maca_stream = dynamic_cast( + infini_train::core::GetDeviceGuardImpl(device.type())->GetStream(device)) + ->maca_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::maca::DispatchMacaFunc, + 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); + } + }, + "MACA Cast"); + + return {dst_tensor}; +} +} // namespace infini_train::kernels::maca + +namespace infini_train::kernels::maca { + +void RegisterCastKernels() { + REGISTER_KERNEL(infini_train::Device::DeviceType::kPrivateUse1, Cast, infini_train::kernels::maca::Cast); +} + +} // namespace infini_train::kernels::maca diff --git a/backends/maca/src/kernels/common/gemm.cc b/backends/maca/src/kernels/common/gemm.cc new file mode 100644 index 0000000..5b6a00b --- /dev/null +++ b/backends/maca/src/kernels/common/gemm.cc @@ -0,0 +1,73 @@ +#include "kernels/common/gemm.h" + +#include + +#include "glog/logging.h" + +#include "common/common_maca.h" +#include "infini_train/include/core/runtime/device_guard.h" +#include "infini_train/include/dispatcher.h" +#include "runtime/maca_runtime_common.h" + +namespace infini_train::kernels::maca { +namespace { + +mcblasOperation_t ToMcblasOperation(GemmTranspose transpose) { + switch (transpose) { + case GemmTranspose::kNoTranspose: + return MCBLAS_OP_N; + case GemmTranspose::kTranspose: + return MCBLAS_OP_T; + } + LOG(FATAL) << "Gemm: unsupported transpose flag " << static_cast(transpose); + return MCBLAS_OP_N; +} + +macaDataType ToMacaDataType(DataType dtype) { + switch (dtype) { + case DataType::kFLOAT32: + return MACA_R_32F; + case DataType::kBFLOAT16: + return MACA_R_16BF; + default: + LOG(FATAL) << "Gemm: unsupported DataType " << static_cast(dtype); + return MACA_R_32F; + } +} + +} // namespace + +void Gemm(Device device, GemmParams params) { + const mcblasHandle_t blas_handle + = dynamic_cast(core::GetDeviceGuardImpl(device.type())->GetBlasHandle(device)) + ->mcblas_handle(); + + if (params.batch_count == 1) { + DCHECK_EQ(params.stride_a, 0LL); + DCHECK_EQ(params.stride_b, 0LL); + DCHECK_EQ(params.stride_c, 0LL); + } + + const mcblasOperation_t trans_a = ToMcblasOperation(params.trans_a); + const mcblasOperation_t trans_b = ToMcblasOperation(params.trans_b); + const macaDataType type_a = ToMacaDataType(params.input_dtype); + const macaDataType type_b = ToMacaDataType(params.input_dtype); + const macaDataType type_c = ToMacaDataType(params.output_dtype); + + if (params.batch_count == 1) { + MCBLAS_CHECK(mcblasGemmEx(blas_handle, trans_a, trans_b, params.m, params.n, params.k, ¶ms.alpha, params.A, + type_a, params.lda, params.B, type_b, params.ldb, ¶ms.beta, params.C, type_c, + params.ldc, MCBLAS_COMPUTE_32F, MCBLAS_GEMM_DEFAULT)); + } else { + MCBLAS_CHECK(mcblasGemmStridedBatchedEx( + blas_handle, trans_a, trans_b, params.m, params.n, params.k, ¶ms.alpha, params.A, type_a, params.lda, + params.stride_a, params.B, type_b, params.ldb, params.stride_b, ¶ms.beta, params.C, type_c, params.ldc, + params.stride_c, params.batch_count, MCBLAS_COMPUTE_32F, MCBLAS_GEMM_DEFAULT)); + } +} + +void RegisterGemmKernels() { + REGISTER_KERNEL(Device::DeviceType::kPrivateUse1, Gemm, infini_train::kernels::maca::Gemm) +} + +} // namespace infini_train::kernels::maca diff --git a/backends/maca/src/kernels/common/gemm.h b/backends/maca/src/kernels/common/gemm.h new file mode 100644 index 0000000..cc9d0a9 --- /dev/null +++ b/backends/maca/src/kernels/common/gemm.h @@ -0,0 +1,9 @@ +#pragma once + +#include "infini_train/src/kernels/common/gemm.h" + +namespace infini_train::kernels::maca { + +void Gemm(Device device, GemmParams params); + +} // namespace infini_train::kernels::maca diff --git a/backends/maca/src/kernels/concat.maca b/backends/maca/src/kernels/concat.maca new file mode 100644 index 0000000..b391012 --- /dev/null +++ b/backends/maca/src/kernels/concat.maca @@ -0,0 +1,255 @@ +#include +#include +#include +#include +#include +#include +#include +#include + +#include "glog/logging.h" + +#include "common/common_maca.h" +#include "infini_train/include/core/runtime/device_guard.h" +#include "infini_train/include/dispatcher.h" +#include "infini_train/include/tensor.h" + +#include "runtime/maca_dispatch.h" +#include "runtime/maca_runtime_common.h" + +namespace infini_train::kernels::maca { +__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)) + ->maca_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::maca::DispatchMacaFunc( + 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; + + MACA_CHECK(mcMallocAsync(reinterpret_cast(&device_input_ptrs), sizeof(T *) * num_inputs, stream)); + MACA_CHECK(mcMemcpyAsync(device_input_ptrs, host_input_ptrs.data(), sizeof(T *) * num_inputs, + mcMemcpyHostToDevice, stream)); + + MACA_CHECK( + mcMallocAsync(reinterpret_cast(&device_offsets), sizeof(int64_t) * (num_inputs + 1), stream)); + MACA_CHECK(mcMemcpyAsync(device_offsets, host_offsets.data(), sizeof(int64_t) * (num_inputs + 1), + mcMemcpyHostToDevice, stream)); + + ConcatForwardKernel<<>>( + device_input_ptrs, static_cast(output->DataPtr()), device_offsets, N, D, num_inputs, K_total); + + MACA_CHECK(mcFreeAsync(device_input_ptrs, stream)); + MACA_CHECK(mcFreeAsync(device_offsets, stream)); + }, + "MACA 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); + // ConcatBackwardKernel maps every grad_output element to exactly one grad tensor element; no Fill is needed. + 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)) + ->maca_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::maca::DispatchMacaFunc( + 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; + + MACA_CHECK(mcMallocAsync(reinterpret_cast(&device_ptrs), sizeof(T *) * num_inputs, stream)); + MACA_CHECK( + mcMemcpyAsync(device_ptrs, host_ptrs.data(), sizeof(T *) * num_inputs, mcMemcpyHostToDevice, stream)); + + MACA_CHECK( + mcMallocAsync(reinterpret_cast(&device_offsets), sizeof(int64_t) * (num_inputs + 1), stream)); + MACA_CHECK(mcMemcpyAsync(device_offsets, host_offsets.data(), sizeof(int64_t) * (num_inputs + 1), + mcMemcpyHostToDevice, stream)); + + ConcatBackwardKernel<<>>( + static_cast(grad_output->DataPtr()), device_ptrs, device_offsets, N, D, num_inputs, K_total); + + MACA_CHECK(mcFreeAsync(device_ptrs, stream)); + MACA_CHECK(mcFreeAsync(device_offsets, stream)); + }, + "MACA ConcatBackward"); + + return grads; +} +} // namespace infini_train::kernels::maca + +namespace infini_train::kernels::maca { + +void RegisterConcatKernels() { + REGISTER_KERNEL(infini_train::Device::DeviceType::kPrivateUse1, ConcatForward, + infini_train::kernels::maca::ConcatForward) + REGISTER_KERNEL(infini_train::Device::DeviceType::kPrivateUse1, ConcatBackward, + infini_train::kernels::maca::ConcatBackward) +} + +} // namespace infini_train::kernels::maca diff --git a/backends/maca/src/kernels/cross_entropy.maca b/backends/maca/src/kernels/cross_entropy.maca new file mode 100644 index 0000000..b2943ff --- /dev/null +++ b/backends/maca/src/kernels/cross_entropy.maca @@ -0,0 +1,236 @@ +#include +#include +#include +#include +#include +#include +#include +#include + +#include +#include + +#include "infini_train/include/core/runtime/device_guard.h" +#include "infini_train/include/dispatcher.h" +#include "infini_train/include/tensor.h" + +#include "common/common_maca.h" +#include "common/cub_compat.cuh" +#include "common/kernel_helper.cuh" +#include "runtime/maca_dispatch.h" +#include "runtime/maca_runtime_common.h" + +namespace infini_train::kernels::maca { +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 cub::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::maca::Cast(input_ptr[base + i])); + } + const float block_max = cub::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::maca::Cast(input_ptr[base + i]) - shared.max_logit); + } + const float block_sum = cub::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::maca::Cast(input_ptr[base + common::maca::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 &maca_stream = dynamic_cast( + infini_train::core::GetDeviceGuardImpl(device.type())->GetStream(device)) + ->maca_stream(); + + return core::maca::DispatchMacaFunc, + 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::maca::Cast(val); }) + / bs; + + return std::make_shared(loss->To(input->GetDevice())); + }, + "MACA 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 cub::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::maca::Cast(input_ptr[idx_base + i])); + } + const float block_max = cub::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::maca::Cast(input_ptr[idx_base + i]) - shared.max_logit); + } + + const float block_sum = cub::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::maca::Cast(input_ptr[global_idx]) - shared.max_logit); + input_grad_ptr[global_idx] = common::maca::Cast((exp_val * scale - (i == target)) * inv_bs + * common::maca::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 &maca_stream = dynamic_cast( + infini_train::core::GetDeviceGuardImpl(device.type())->GetStream(device)) + ->maca_stream(); + + core::maca::DispatchMacaFunc, + DataTypeList>( + {target->Dtype(), input_casted->Dtype()}, + [=]() { + // One sample block writes all of its num_classes gradient elements; no Fill is needed. + 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); + }, + "MACA CrossEntropyBackward"); + + return {grad_input}; +} +} // namespace infini_train::kernels::maca + +namespace infini_train::kernels::maca { + +void RegisterCrossEntropyKernels() { + REGISTER_KERNEL(infini_train::Device::DeviceType::kPrivateUse1, CrossEntropyForward, + infini_train::kernels::maca::CrossEntropyForward) + REGISTER_KERNEL(infini_train::Device::DeviceType::kPrivateUse1, CrossEntropyBackward, + infini_train::kernels::maca::CrossEntropyBackward) +} + +} // namespace infini_train::kernels::maca diff --git a/backends/maca/src/kernels/elementwise.maca b/backends/maca/src/kernels/elementwise.maca new file mode 100644 index 0000000..fbc5155 --- /dev/null +++ b/backends/maca/src/kernels/elementwise.maca @@ -0,0 +1,1311 @@ +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include + +#include + +#include "infini_train/include/core/runtime/device_guard.h" +#include "infini_train/include/dispatcher.h" +#include "infini_train/include/tensor.h" + +#include "common/common_maca.h" +#include "common/kernel_helper.cuh" +#include "runtime/maca_dispatch.h" +#include "runtime/maca_runtime_common.h" + +namespace infini_train::kernels::maca { +namespace { +using namespace infini_train::common::maca; +constexpr int kWarpSize = 32; + +// 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 CUDA kernel parameter limit. +constexpr int kMaxBroadcastDims = 8; + +// POD metadata for broadcast kernels. Passed by value into __global__ kernels so the data +// lives in CUDA kernel parameter memory (constant cache) instead of being uploaded via a +// per-call mcMallocAsync + mcMemcpyAsync 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; +} + +inline dim3 ChooseBlockDims(size_t num_elements) { return dim3(ChooseBlockSize(num_elements)); } + +// 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(); + dim3 block_dims = ChooseBlockDims(num_elements); + 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 &maca_stream = dynamic_cast( + infini_train::core::GetDeviceGuardImpl(device.type())->GetStream(device)) + ->maca_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 mcMalloc/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 = ChooseBlockDims(num_elements); + 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 mcMallocAsync/mcMemcpyAsync/mcFreeAsync 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::maca::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::maca::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::maca::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::maca::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, mcStream_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; + MACA_CHECK(mcMallocAsync(reinterpret_cast(&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); + MACA_CHECK(mcGetLastError()); + + // Pass 2: choose between 1D and 2D reductions depending on workload shape. + int dev = 0; + int sm_count = 0; + MACA_CHECK(mcGetDevice(&dev)); + MACA_CHECK(mcDeviceGetAttribute(&sm_count, mcDeviceAttributeMultiProcessorCount, 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); + MACA_CHECK(mcGetLastError()); + } 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; + MACA_CHECK( + mcMallocAsync(reinterpret_cast(&outB_accum), static_cast(K) * sizeof(float), stream)); + MACA_CHECK(mcMemsetAsync(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); + MACA_CHECK(mcGetLastError()); + + // Convert accumulated floats back to the target dtype. + const dim3 cgrid((K + RED_THREADS - 1) / RED_THREADS); + CastFloatToTBhist<<>>(outB_accum, outB, K); + MACA_CHECK(mcGetLastError()); + + MACA_CHECK(mcFreeAsync(outB_accum, stream)); + } + + MACA_CHECK(mcFreeAsync(work, stream)); +} + +// 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 lane_id = tid % kWarpSize; + const int logical_warp_id = tid / kWarpSize; + + using WarpReduce = cub::WarpReduce; + auto *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::maca::Cast(Mul(grad_output[idx], fn_b(a_val, b_val))); + } + + using WarpMask = decltype(__ballot_sync(~uint64_t{0}, true)); + const WarpMask full_mask = ~WarpMask{0}; + const WarpMask physical_active_mask = __ballot_sync(full_mask, in_bounds); + const int physical_lane = tid % warpSize; + const int logical_base = (physical_lane / kWarpSize) * kWarpSize; + const WarpMask logical_lane_mask = static_cast(uint64_t{0xffffffff} << logical_base); + const WarpMask active_mask = physical_active_mask & logical_lane_mask; + if (active_mask == 0) { + return; + } + + const unsigned logical_active_mask = static_cast(static_cast(active_mask) >> logical_base); + const int leader = __ffs(logical_active_mask) - 1; + const int64_t common_offset = __shfl_sync(active_mask, b_offset, leader, kWarpSize); + + bool warp_uniform = true; + for (int i = 0; i < kWarpSize; ++i) { + if (!(logical_active_mask & (unsigned{1} << i))) { + continue; + } + const int64_t offset_i = __shfl_sync(active_mask, b_offset, i, kWarpSize); + if (offset_i != common_offset) { + warp_uniform = false; + break; + } + } + + if (warp_uniform) { + const float reduced = WarpReduce(temp_storage[logical_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 + atomicAdd(&output_b[common_offset], common::maca::Cast(reduced)); + } + } else if (in_bounds) { + // FIXME(lzm): atomicAdd is much slower for bf16 and __half compared to float, needs further optimization + atomicAdd(&output_b[b_offset], common::maca::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::maca::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::maca::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 &maca_stream = dynamic_cast( + infini_train::core::GetDeviceGuardImpl(device.type())->GetStream(device)) + ->maca_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)) + ->maca_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 mcMalloc/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 = ChooseBlockDims(num_elements); + 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 mcMallocAsync/mcMemcpyAsync/mcFreeAsync 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 block_threads = static_cast(block.x); + const int num_warps = CEIL_DIV(block_threads, kWarpSize); + const size_t smem_size = num_warps * sizeof(cub::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 block_threads = static_cast(block.x); + const int padded_block = block_threads + block_threads / 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(unary_fn, output, input);), DataType::kFLOAT32) + DISPATCH_CASE(WRAP(LaunchForward<__maca_bfloat16>(unary_fn, output, input);), DataType::kBFLOAT16) + DISPATCH_CASE(WRAP(LaunchForward(unary_fn, output, input);), DataType::kINT64) + default: + LOG_LOC(FATAL, "MACA 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(unary_fn, output, grad_output_promoted, a_promoted); }), + DataType::kFLOAT32) + DISPATCH_CASE(WRAP({ LaunchBackward<__maca_bfloat16>(unary_fn, output, grad_output_promoted, a_promoted); }), + DataType::kBFLOAT16) + DISPATCH_CASE(WRAP({ LaunchBackward(unary_fn, output, grad_output_promoted, a_promoted); }), + DataType::kINT64) + default: + LOG_LOC(FATAL, "MACA 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(binary_fn, output, a_promoted, b_promoted);), DataType::kFLOAT32) + DISPATCH_CASE(WRAP(LaunchForward<__maca_bfloat16>(binary_fn, output, a_promoted, b_promoted);), + DataType::kBFLOAT16) + DISPATCH_CASE(WRAP(LaunchForward(binary_fn, output, a_promoted, b_promoted);), DataType::kINT64) + default: + LOG_LOC(FATAL, "MACA 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(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); + grad_b->Fill(0); + } + LaunchBackward<__maca_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); + // grad_b->Fill(0); + // LaunchBackward(fn_a, fn_b, grad_a, grad_b, a_dims, b_dims, grad_output, a, b); + // }), + // DataType::kINT64) + default: + LOG_LOC(FATAL, "MACA 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::maca::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_NUMERIC_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_NUMERIC_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_NUMERIC_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_NUMERIC_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_NUMERIC_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_NUMERIC_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_NUMERIC_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_NUMERIC_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_NUMERIC_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_NUMERIC_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_NUMERIC_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_NUMERIC_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::maca::Cast(1); }); + , INFINI_ALL_NUMERIC_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_NUMERIC_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, "MACA 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, "MACA 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::maca + +namespace infini_train::kernels::maca { + +void RegisterElementwiseKernels() { + REGISTER_KERNEL(infini_train::Device::DeviceType::kPrivateUse1, NegForward, infini_train::kernels::maca::NegForward) + REGISTER_KERNEL(infini_train::Device::DeviceType::kPrivateUse1, NegBackward, + infini_train::kernels::maca::NegBackward) + REGISTER_KERNEL(infini_train::Device::DeviceType::kPrivateUse1, ReciprocalForward, + infini_train::kernels::maca::ReciprocalForward) + REGISTER_KERNEL(infini_train::Device::DeviceType::kPrivateUse1, ReciprocalBackward, + infini_train::kernels::maca::ReciprocalBackward) + REGISTER_KERNEL(infini_train::Device::DeviceType::kPrivateUse1, SinForward, infini_train::kernels::maca::SinForward) + REGISTER_KERNEL(infini_train::Device::DeviceType::kPrivateUse1, SinBackward, + infini_train::kernels::maca::SinBackward) + REGISTER_KERNEL(infini_train::Device::DeviceType::kPrivateUse1, CosForward, infini_train::kernels::maca::CosForward) + REGISTER_KERNEL(infini_train::Device::DeviceType::kPrivateUse1, CosBackward, + infini_train::kernels::maca::CosBackward) + REGISTER_KERNEL(infini_train::Device::DeviceType::kPrivateUse1, TanhForward, + infini_train::kernels::maca::TanhForward) + REGISTER_KERNEL(infini_train::Device::DeviceType::kPrivateUse1, TanhBackward, + infini_train::kernels::maca::TanhBackward) + REGISTER_KERNEL(infini_train::Device::DeviceType::kPrivateUse1, PowForward, infini_train::kernels::maca::PowForward) + REGISTER_KERNEL(infini_train::Device::DeviceType::kPrivateUse1, PowBackward, + infini_train::kernels::maca::PowBackward) + REGISTER_KERNEL(infini_train::Device::DeviceType::kPrivateUse1, RsqrtForward, + infini_train::kernels::maca::RsqrtForward) + REGISTER_KERNEL(infini_train::Device::DeviceType::kPrivateUse1, RsqrtBackward, + infini_train::kernels::maca::RsqrtBackward) + REGISTER_KERNEL(infini_train::Device::DeviceType::kPrivateUse1, ExpForward, infini_train::kernels::maca::ExpForward) + REGISTER_KERNEL(infini_train::Device::DeviceType::kPrivateUse1, ExpBackward, + infini_train::kernels::maca::ExpBackward) + REGISTER_KERNEL(infini_train::Device::DeviceType::kPrivateUse1, LogForward, infini_train::kernels::maca::LogForward) + REGISTER_KERNEL(infini_train::Device::DeviceType::kPrivateUse1, LogBackward, + infini_train::kernels::maca::LogBackward) + REGISTER_KERNEL(infini_train::Device::DeviceType::kPrivateUse1, EqualsForward, + infini_train::kernels::maca::EqualsForward) + REGISTER_KERNEL(infini_train::Device::DeviceType::kPrivateUse1, EqualsScalarForward, + infini_train::kernels::maca::EqualsScalarForward) + REGISTER_KERNEL(infini_train::Device::DeviceType::kPrivateUse1, LtForward, infini_train::kernels::maca::LtForward) + REGISTER_KERNEL(infini_train::Device::DeviceType::kPrivateUse1, LtScalarForward, + infini_train::kernels::maca::LtScalarForward) + REGISTER_KERNEL(infini_train::Device::DeviceType::kPrivateUse1, LeForward, infini_train::kernels::maca::LeForward) + REGISTER_KERNEL(infini_train::Device::DeviceType::kPrivateUse1, LeScalarForward, + infini_train::kernels::maca::LeScalarForward) + REGISTER_KERNEL(infini_train::Device::DeviceType::kPrivateUse1, GtForward, infini_train::kernels::maca::GtForward) + REGISTER_KERNEL(infini_train::Device::DeviceType::kPrivateUse1, GtScalarForward, + infini_train::kernels::maca::GtScalarForward) + REGISTER_KERNEL(infini_train::Device::DeviceType::kPrivateUse1, GeForward, infini_train::kernels::maca::GeForward) + REGISTER_KERNEL(infini_train::Device::DeviceType::kPrivateUse1, GeScalarForward, + infini_train::kernels::maca::GeScalarForward) + REGISTER_KERNEL(infini_train::Device::DeviceType::kPrivateUse1, OrForward, infini_train::kernels::maca::OrForward) + REGISTER_KERNEL(infini_train::Device::DeviceType::kPrivateUse1, AndForward, infini_train::kernels::maca::AndForward) + REGISTER_KERNEL(infini_train::Device::DeviceType::kPrivateUse1, AddForward, infini_train::kernels::maca::AddForward) + REGISTER_KERNEL(infini_train::Device::DeviceType::kPrivateUse1, AddBackward, + infini_train::kernels::maca::AddBackward) + REGISTER_KERNEL(infini_train::Device::DeviceType::kPrivateUse1, AddScalarForward, + infini_train::kernels::maca::AddScalarForward) + REGISTER_KERNEL(infini_train::Device::DeviceType::kPrivateUse1, AddScalarBackward, + infini_train::kernels::maca::AddScalarBackward) + REGISTER_KERNEL(infini_train::Device::DeviceType::kPrivateUse1, SubForward, infini_train::kernels::maca::SubForward) + REGISTER_KERNEL(infini_train::Device::DeviceType::kPrivateUse1, SubBackward, + infini_train::kernels::maca::SubBackward) + REGISTER_KERNEL(infini_train::Device::DeviceType::kPrivateUse1, MulForward, infini_train::kernels::maca::MulForward) + REGISTER_KERNEL(infini_train::Device::DeviceType::kPrivateUse1, MulBackward, + infini_train::kernels::maca::MulBackward) + REGISTER_KERNEL(infini_train::Device::DeviceType::kPrivateUse1, MulScalarForward, + infini_train::kernels::maca::MulScalarForward) + REGISTER_KERNEL(infini_train::Device::DeviceType::kPrivateUse1, MulScalarBackward, + infini_train::kernels::maca::MulScalarBackward) + REGISTER_KERNEL(infini_train::Device::DeviceType::kPrivateUse1, DivForward, infini_train::kernels::maca::DivForward) + REGISTER_KERNEL(infini_train::Device::DeviceType::kPrivateUse1, DivBackward, + infini_train::kernels::maca::DivBackward) + REGISTER_KERNEL(infini_train::Device::DeviceType::kPrivateUse1, SigmoidForward, + infini_train::kernels::maca::SigmoidForward) + REGISTER_KERNEL(infini_train::Device::DeviceType::kPrivateUse1, SigmoidBackward, + infini_train::kernels::maca::SigmoidBackward) +} + +} // namespace infini_train::kernels::maca diff --git a/backends/maca/src/kernels/embedding.maca b/backends/maca/src/kernels/embedding.maca new file mode 100644 index 0000000..0000acf --- /dev/null +++ b/backends/maca/src/kernels/embedding.maca @@ -0,0 +1,136 @@ +#include +#include +#include +#include +#include + +#include "common/common_maca.h" +#include "infini_train/include/core/runtime/device_guard.h" +#include "infini_train/include/dispatcher.h" +#include "infini_train/include/tensor.h" + +#include "runtime/maca_dispatch.h" +#include "runtime/maca_runtime_common.h" + +namespace infini_train::kernels::maca { + +template +__global__ void EmbeddingForwardKernel(const int64_t *input, T *output, const T *weight, size_t num_output_elements, + size_t embed_dim, int64_t vocab_size) { + 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_output_elements; + idx += grid_stride) { + const size_t token_idx = idx / embed_dim; + const size_t channel = idx % embed_dim; + const int64_t token_id = input[token_idx]; + if (token_id < 0 || token_id >= vocab_size) { + continue; + } + output[idx] = weight[static_cast(token_id) * embed_dim + channel]; + } +} + +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 &maca_stream = dynamic_cast( + infini_train::core::GetDeviceGuardImpl(device.type())->GetStream(device)) + ->maca_stream(); + + const int64_t vocab_size = weight->Dims()[0]; + const size_t embed_dim = static_cast(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()); + const size_t num_output_elements = output->NumElements(); + if (num_output_elements == 0) { + return output; + } + + constexpr int threads_per_block = 256; + const int num_blocks = static_cast( + std::min((num_output_elements + threads_per_block - 1) / threads_per_block, static_cast(65535))); + + core::maca::DispatchMacaFunc( + dtype, + [=]() { + EmbeddingForwardKernel<<>>( + static_cast(input->DataPtr()), static_cast(output->DataPtr()), + static_cast(weight->DataPtr()), num_output_elements, embed_dim, vocab_size); + }, + "MACA EmbeddingForward"); + + return output; +} + +template +__global__ void EmbeddingBackwardKernel(const int64_t *input_ptr, const T *grad_output_ptr, T *grad_weight_ptr, + size_t num_tokens, size_t embedding_dim, int64_t vocab_size) { + 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_tokens; + idx += grid_stride) { + const int64_t token_id = input_ptr[idx]; + if (token_id < 0 || token_id >= vocab_size) { + continue; + } + + const size_t weight_offset = static_cast(token_id) * embedding_dim; + const size_t grad_output_offset = idx * embedding_dim; + for (size_t j = 0; j < embedding_dim; ++j) { + atomicAdd(&grad_weight_ptr[weight_offset + j], grad_output_ptr[grad_output_offset + 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 &maca_stream = dynamic_cast( + infini_train::core::GetDeviceGuardImpl(device.type())->GetStream(device)) + ->maca_stream(); + + const int64_t vocab_size = weight_dims[0]; + const size_t embedding_dim = static_cast(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 size_t num_tokens = input->NumElements(); + constexpr int threads_per_block = 256; + const int num_blocks = static_cast( + std::min((num_tokens + threads_per_block - 1) / threads_per_block, static_cast(65535))); + + core::maca::DispatchMacaFunc( + dtype, + [=]() { + grad_weight->Fill(0); + if (num_tokens > 0) { + EmbeddingBackwardKernel<<>>( + static_cast(input->DataPtr()), static_cast(grad_output->DataPtr()), + static_cast(grad_weight->DataPtr()), num_tokens, embedding_dim, vocab_size); + } + }, + "MACA EmbeddingBackward"); + + return grad_weight; +} +} // namespace infini_train::kernels::maca + +namespace infini_train::kernels::maca { + +void RegisterEmbeddingKernels() { + REGISTER_KERNEL(infini_train::Device::DeviceType::kPrivateUse1, EmbeddingForward, + infini_train::kernels::maca::EmbeddingForward) + REGISTER_KERNEL(infini_train::Device::DeviceType::kPrivateUse1, EmbeddingBackward, + infini_train::kernels::maca::EmbeddingBackward) +} + +} // namespace infini_train::kernels::maca diff --git a/backends/maca/src/kernels/fill.maca b/backends/maca/src/kernels/fill.maca new file mode 100644 index 0000000..b815d91 --- /dev/null +++ b/backends/maca/src/kernels/fill.maca @@ -0,0 +1,59 @@ +#include +#include +#include + +#include "common/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 "runtime/maca_dispatch.h" +#include "runtime/maca_runtime_common.h" + +namespace infini_train::kernels::maca { + +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 &maca_stream = dynamic_cast( + infini_train::core::GetDeviceGuardImpl(device.type())->GetStream(device)) + ->maca_stream(); + + core::maca::DispatchMacaFunc( + tensor->Dtype(), + [=]() { + // Scalar::to relies on static_cast, which is ambiguous when T is a + // MACA native half/bf16 type constructed from integer scalars. Route + // half/bf16 through float via common::maca::Cast to guarantee a + // single-candidate conversion path. + T casted_value; + if constexpr (std::is_same_v || std::is_same_v) { + casted_value = common::maca::Cast(scalar.to()); + } else { + casted_value = scalar.to(); + } + FillKernel<<>>(static_cast(tensor->DataPtr()), + casted_value, tensor->NumElements()); + }, + "MACA Fill"); +} +} // namespace infini_train::kernels::maca + +namespace infini_train::kernels::maca { + +void RegisterFillKernels() { + REGISTER_KERNEL(infini_train::Device::DeviceType::kPrivateUse1, Fill, infini_train::kernels::maca::Fill); +} + +} // namespace infini_train::kernels::maca diff --git a/backends/maca/src/kernels/gather.maca b/backends/maca/src/kernels/gather.maca new file mode 100644 index 0000000..85e2699 --- /dev/null +++ b/backends/maca/src/kernels/gather.maca @@ -0,0 +1,237 @@ +#include +#include +#include +#include +#include +#include + +#include "glog/logging.h" + +#include "common/common_maca.h" +#include "infini_train/include/core/runtime/device_guard.h" +#include "infini_train/include/dispatcher.h" +#include "infini_train/include/tensor.h" + +#include "runtime/maca_dispatch.h" +#include "runtime/maca_runtime_common.h" + +namespace infini_train::kernels::maca { +template +__global__ void GatherForwardKernel(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 GatherForward(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)) + ->maca_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; + MACA_CHECK(mcMallocAsync(reinterpret_cast(&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; + + MACA_CHECK(mcMemcpyAsync(out_dims_dev, idx_dims.data(), num_dims * sizeof(int64_t), mcMemcpyHostToDevice, stream)); + MACA_CHECK( + mcMemcpyAsync(in_strides_dev, in_strides.data(), num_dims * sizeof(int64_t), mcMemcpyHostToDevice, stream)); + MACA_CHECK( + mcMemcpyAsync(out_strides_dev, out_strides.data(), num_dims * sizeof(int64_t), mcMemcpyHostToDevice, stream)); + + const int threads = 256; + const int blocks = (total_elements + threads - 1) / threads; + + core::maca::DispatchMacaFunc( + dtype, + [=]() { + GatherForwardKernel<<>>( + 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); + }, + "MACA GatherForward"); + + MACA_CHECK(mcFreeAsync(dev_buf, stream)); + return out; +} + +template +__global__ void GatherBackwardKernel(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]; + } + } + atomicAdd(&grad_input[in_linear], grad_output[out_idx]); +} + +std::shared_ptr GatherBackward(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_LE(idx_dims[d], in_dims[d]) + << "index.size(" << d << ") must be <= input.size(" << d << ") on non-gather dims"; + } + + auto dtype = grad_output->Dtype(); + auto grad_input = std::make_shared(in_dims, dtype, grad_output->GetDevice()); + core::maca::DispatchMacaFunc( + dtype, [=]() { grad_input->Fill(0); }, "MACA GatherBackwardZero"); + + 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)) + ->maca_stream(); + + MACA_CHECK(mcMallocAsync(reinterpret_cast(&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; + + MACA_CHECK(mcMemcpyAsync(out_dims_dev, idx_dims.data(), n_out * sizeof(int64_t), mcMemcpyHostToDevice, stream)); + MACA_CHECK( + mcMemcpyAsync(in_strides_dev, in_strides.data(), n_in_strides * sizeof(int64_t), mcMemcpyHostToDevice, stream)); + MACA_CHECK(mcMemcpyAsync(out_strides_dev, out_strides.data(), n_out_strides * sizeof(int64_t), mcMemcpyHostToDevice, + stream)); + + const int threads = 256; + const int blocks = (int)((total_elements + threads - 1) / threads); + + core::maca::DispatchMacaFunc( + dtype, + [=]() { + GatherBackwardKernel<<>>( + 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); + }, + "MACA GatherBackward"); + + MACA_CHECK(mcFreeAsync(dev_buf, stream)); + return grad_input; +} + +} // namespace infini_train::kernels::maca + +namespace infini_train::kernels::maca { + +void RegisterGatherKernels() { + REGISTER_KERNEL(infini_train::Device::DeviceType::kPrivateUse1, GatherForward, + infini_train::kernels::maca::GatherForward) + REGISTER_KERNEL(infini_train::Device::DeviceType::kPrivateUse1, GatherBackward, + infini_train::kernels::maca::GatherBackward) +} + +} // namespace infini_train::kernels::maca diff --git a/backends/maca/src/kernels/layernorm.maca b/backends/maca/src/kernels/layernorm.maca new file mode 100644 index 0000000..86f3b79 --- /dev/null +++ b/backends/maca/src/kernels/layernorm.maca @@ -0,0 +1,248 @@ +#include +#include +#include +#include +#include + +#include + +#include "common/common_maca.h" +#include "common/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 "runtime/maca_dispatch.h" +#include "runtime/maca_runtime_common.h" + +namespace infini_train::kernels::maca { + +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 = cub::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::maca::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::maca::Cast(x[i]) - shared_mean) * shared_rstd; + y[i] = common::maca::Cast(norm * common::maca::Cast(weight[i]) + common::maca::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 &maca_stream = dynamic_cast( + infini_train::core::GetDeviceGuardImpl(device.type())->GetStream(device)) + ->maca_stream(); + + core::maca::DispatchMacaFunc( + dtype, + [=]() { + // Each token block writes its mean and rstd exactly once; no Fill is needed. + 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); + }, + "MACA LayerNormForward"); + + return {output, mean, rstd}; +} + +template +__global__ void LayerNormInputGradKernel(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, int embed_dim) { + using BlockReduce = cub::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::maca::Cast(weight[i]) * common::maca::Cast(grad_output_ptr[i]); + float norm = (common::maca::Cast(input_ptr[i]) - mean_val) * rstd_val; + dnorm_mean += dnorm; + dnorm_norm_mean += dnorm * norm; + } + + 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; + shared_mean = mean_d; + shared_norm = norm_d; + } + __syncthreads(); + + for (int i = tid; i < embed_dim; i += BLOCK_SIZE) { + float norm = (common::maca::Cast(input_ptr[i]) - mean_val) * rstd_val; + float grad_output_val = common::maca::Cast(grad_output_ptr[i]); + + grad_input_ptr[i] = common::maca::Cast( + (common::maca::Cast(weight[i]) * grad_output_val - shared_mean - norm * shared_norm) * rstd_val); + } +} + +template +__global__ void LayerNormParameterGradKernel(const T *__restrict__ input, const T *__restrict__ grad_output, + const float *__restrict__ mean, const float *__restrict__ rstd, + T *__restrict__ grad_weight, T *__restrict__ grad_bias, int64_t num_tokens, + int embed_dim) { + using BlockReduce = cub::BlockReduce; + __shared__ typename BlockReduce::TempStorage temp_storage_weight; + __shared__ typename BlockReduce::TempStorage temp_storage_bias; + + int feature_idx = blockIdx.x; + float grad_weight_sum = 0.0f; + float grad_bias_sum = 0.0f; + + for (int64_t token_idx = threadIdx.x; token_idx < num_tokens; token_idx += BLOCK_SIZE) { + int64_t offset = token_idx * embed_dim + feature_idx; + float grad_output_val = common::maca::Cast(grad_output[offset]); + float norm = (common::maca::Cast(input[offset]) - mean[token_idx]) * rstd[token_idx]; + grad_weight_sum += grad_output_val * norm; + grad_bias_sum += grad_output_val; + } + + float grad_weight_reduced = BlockReduce(temp_storage_weight).Sum(grad_weight_sum); + float grad_bias_reduced = BlockReduce(temp_storage_bias).Sum(grad_bias_sum); + + if (threadIdx.x == 0) { + grad_weight[feature_idx] = common::maca::Cast(grad_weight_reduced); + grad_bias[feature_idx] = common::maca::Cast(grad_bias_reduced); + } +} + +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) { + 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]; + const int64_t num_tokens = static_cast(batch_size) * max_seqlen; + + 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 &maca_stream = dynamic_cast( + infini_train::core::GetDeviceGuardImpl(device.type())->GetStream(device)) + ->maca_stream(); + core::maca::DispatchMacaFunc( + dtype, + [=]() { + grad_weight->Fill(0); + grad_bias->Fill(0); + LayerNormInputGradKernel<<>>( + 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()), embed_dim); + LayerNormParameterGradKernel<<>>( + static_cast(input->DataPtr()), static_cast(grad_output->DataPtr()), + static_cast(mean->DataPtr()), static_cast(rstd->DataPtr()), + static_cast(grad_weight->DataPtr()), static_cast(grad_bias->DataPtr()), num_tokens, + embed_dim); + }, + "MACA LayerNormBackward"); + + return {grad_input, grad_weight, grad_bias}; +} +} // namespace infini_train::kernels::maca + +namespace infini_train::kernels::maca { + +void RegisterLayerNormKernels() { + REGISTER_KERNEL(infini_train::Device::DeviceType::kPrivateUse1, LayerNormForward, + infini_train::kernels::maca::LayerNormForward) + REGISTER_KERNEL(infini_train::Device::DeviceType::kPrivateUse1, LayerNormBackward, + infini_train::kernels::maca::LayerNormBackward) +} + +} // namespace infini_train::kernels::maca diff --git a/backends/maca/src/kernels/linear.maca b/backends/maca/src/kernels/linear.maca new file mode 100644 index 0000000..54bfa4f --- /dev/null +++ b/backends/maca/src/kernels/linear.maca @@ -0,0 +1,544 @@ +#include +#include +#include +#include +#include + +#include + +#include "infini_train/include/autograd/linear.h" +#include "infini_train/include/core/runtime/device_guard.h" +#include "infini_train/include/dispatcher.h" +#include "infini_train/include/tensor.h" + +#include "common/common_maca.h" +#include "common/kernel_helper.cuh" +#include "kernels/common/gemm.h" +#include "runtime/maca_dispatch.h" +#include "runtime/maca_runtime_common.h" + +namespace infini_train::kernels::maca { + +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(), int64_t{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(); + // mcBLAS is column-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] + Dispatcher::Instance().Call({device.type(), "Gemm"}, device, + GemmParams{ + .trans_a = GemmTranspose::kNoTranspose, + .trans_b = GemmTranspose::kNoTranspose, + .m = static_cast(n), + .n = static_cast(m), + .k = static_cast(k), + .A = other->DataPtr(), + .lda = static_cast(n), + .B = input->DataPtr(), + .ldb = static_cast(k), + .C = output->DataPtr(), + .ldc = static_cast(n), + .alpha = 1.0f, + .beta = 0.0f, + .batch_count = static_cast(bs), + .stride_a = bs > 1 ? n * k : 0, + .stride_b = bs > 1 ? k * m : 0, + .stride_c = bs > 1 ? m * n : 0, + .input_dtype = dtype, + .output_dtype = dtype, + }); + + return output; +} + +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(), int64_t{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)); + + // FIXME(cx): output dtype promotion is a temporary hack; revisit when autograd/autocast is fixed. + 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(); + Dispatcher::Instance().Call({device.type(), "Gemm"}, device, + GemmParams{ + .trans_a = GemmTranspose::kTranspose, + .trans_b = GemmTranspose::kNoTranspose, + .m = static_cast(k), + .n = static_cast(m), + .k = static_cast(n), + .A = other->DataPtr(), + .lda = static_cast(n), + .B = grad_output_promoted->DataPtr(), + .ldb = static_cast(n), + .C = grad_input->DataPtr(), + .ldc = static_cast(k), + .alpha = 1.0f, + .beta = 0.0f, + .batch_count = static_cast(bs), + .stride_a = bs > 1 ? k * n : 0, + .stride_b = bs > 1 ? n * m : 0, + .stride_c = bs > 1 ? m * k : 0, + .input_dtype = compute_dtype, + .output_dtype = output_dtype, + }); + + 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(), int64_t{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)); + + // FIXME(cx): output dtype promotion is a temporary hack; revisit when autograd/autocast is fixed. + 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(); + Dispatcher::Instance().Call({device.type(), "Gemm"}, device, + GemmParams{ + .trans_a = GemmTranspose::kNoTranspose, + .trans_b = GemmTranspose::kTranspose, + .m = static_cast(n), + .n = static_cast(k), + .k = static_cast(m), + .A = grad_output_promoted->DataPtr(), + .lda = static_cast(n), + .B = input->DataPtr(), + .ldb = static_cast(k), + .C = grad_other->DataPtr(), + .ldc = static_cast(n), + .alpha = 1.0f, + .beta = 0.0f, + .batch_count = static_cast(bs), + .stride_a = bs > 1 ? n * m : 0, + .stride_b = bs > 1 ? k * m : 0, + .stride_c = bs > 1 ? n * k : 0, + .input_dtype = compute_dtype, + .output_dtype = output_dtype, + }); + + 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) { + return; + } + int j = idx % out_features; + output[idx] = bias[j]; +} + +// Accumulate FP32 products in FP64 with a fixed order, then round once when writing the result. +constexpr int kLinearBackwardInputTileSize = 16; + +__global__ void LinearBackwardInputFp32Kernel(const float *__restrict__ weight, const float *__restrict__ grad_output, + float *__restrict__ grad_input, int64_t bs, int64_t in_features, + int64_t out_features, bool transpose) { + __shared__ float grad_tile[kLinearBackwardInputTileSize][kLinearBackwardInputTileSize]; + __shared__ float weight_tile[kLinearBackwardInputTileSize][kLinearBackwardInputTileSize]; + + const int64_t row + = static_cast(blockIdx.y) * kLinearBackwardInputTileSize + static_cast(threadIdx.y); + const int64_t in_col + = static_cast(blockIdx.x) * kLinearBackwardInputTileSize + static_cast(threadIdx.x); + double sum = 0.0; + + for (int64_t out_base = 0; out_base < out_features; out_base += kLinearBackwardInputTileSize) { + const int64_t grad_col = out_base + static_cast(threadIdx.x); + grad_tile[threadIdx.y][threadIdx.x] + = row < bs && grad_col < out_features ? grad_output[row * out_features + grad_col] : 0.0f; + + const int64_t weight_row = out_base + static_cast(threadIdx.y); + if (weight_row < out_features && in_col < in_features) { + const int64_t weight_index + = transpose ? weight_row * in_features + in_col : in_col * out_features + weight_row; + weight_tile[threadIdx.y][threadIdx.x] = weight[weight_index]; + } else { + weight_tile[threadIdx.y][threadIdx.x] = 0.0f; + } + __syncthreads(); + + const int64_t remaining = out_features - out_base; + const int tile_width + = remaining < kLinearBackwardInputTileSize ? static_cast(remaining) : kLinearBackwardInputTileSize; +#pragma unroll 1 + for (int k = 0; k < tile_width; ++k) { + sum += static_cast(grad_tile[threadIdx.y][k] * weight_tile[k][threadIdx.x]); + } + __syncthreads(); + } + + if (row < bs && in_col < in_features) { + grad_input[row * in_features + in_col] = static_cast(sum); + } +} + +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(), int64_t{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]); + + // mcBLAS computes the equivalent column-major product: + // 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 &maca_stream = dynamic_cast( + infini_train::core::GetDeviceGuardImpl(device.type())->GetStream(device)) + ->maca_stream(); + + const float beta = bias ? 1.0f : 0.0f; + 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::maca::DispatchMacaFunc( + dtype, + [=]() { + BiasCopyKernel<<>>( + static_cast(output->DataPtr()), static_cast(bias->DataPtr()), bs, out_features); + }, + "MACA LinearForward"); + } + + // - 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] + Dispatcher::Instance().Call( + {device.type(), "Gemm"}, device, + GemmParams{ + .trans_a = transpose ? GemmTranspose::kTranspose : GemmTranspose::kNoTranspose, + .trans_b = GemmTranspose::kNoTranspose, + .m = static_cast(out_features), + .n = static_cast(bs), + .k = static_cast(in_features), + .A = weight->DataPtr(), + .lda = static_cast(transpose ? in_features : out_features), + .B = input->DataPtr(), + .ldb = static_cast(in_features), + .C = output->DataPtr(), + .ldc = static_cast(out_features), + .alpha = 1.0f, + .beta = beta, + .batch_count = 1, + .input_dtype = dtype, + .output_dtype = dtype, + }); + + return output; +} + +template +__global__ void ReduceRowsKernel(const TIn *__restrict__ input, TOut *__restrict__ output, int64_t num_rows, + int64_t num_cols) { + using BlockReduce = cub::BlockReduce; + __shared__ typename BlockReduce::TempStorage temp_storage; + + const int64_t col = blockIdx.x; + float sum = 0.0f; + + for (int64_t row = threadIdx.x; row < num_rows; row += blockDim.x) { + sum += common::maca::Cast(input[row * num_cols + col]); + } + + float reduced = BlockReduce(temp_storage).Sum(sum); + + if (threadIdx.x == 0) { + output[col] = common::maca::Cast(reduced); + } +} + +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) { + CHECK_GE(input_dims.size(), 2); + const int64_t bs + = std::accumulate(input_dims.rbegin() + 1, input_dims.rend(), int64_t{1}, std::multiplies{}); + + auto compute_dtype = weight->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)); + + // FIXME(cx): output dtype promotion is a temporary hack; revisit when autograd/autocast is fixed. + 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(); + const auto &maca_stream = dynamic_cast( + infini_train::core::GetDeviceGuardImpl(device.type())->GetStream(device)) + ->maca_stream(); + + // d_input = d_output * weight (transposed weight layout) or d_output * weight.T. + switch (compute_dtype) { + DISPATCH_CASE(WRAP({ + const dim3 block_dims(kLinearBackwardInputTileSize, kLinearBackwardInputTileSize); + const dim3 grid_dims(CEIL_DIV(in_features, kLinearBackwardInputTileSize), + CEIL_DIV(bs, kLinearBackwardInputTileSize)); + LinearBackwardInputFp32Kernel<<>>( + static_cast(weight->DataPtr()), + static_cast(grad_output_promoted->DataPtr()), + static_cast(grad_input->DataPtr()), bs, in_features, out_features, transpose); + MACA_CHECK(mcGetLastError()); + }), + DataType::kFLOAT32) + DISPATCH_CASE(WRAP({ + Dispatcher::Instance().Call( + {device.type(), "Gemm"}, device, + GemmParams{ + .trans_a = transpose ? GemmTranspose::kNoTranspose : GemmTranspose::kTranspose, + .trans_b = GemmTranspose::kNoTranspose, + .m = static_cast(in_features), + .n = static_cast(bs), + .k = static_cast(out_features), + .A = weight->DataPtr(), + .lda = static_cast(transpose ? in_features : out_features), + .B = grad_output_promoted->DataPtr(), + .ldb = static_cast(out_features), + .C = grad_input->DataPtr(), + .ldc = static_cast(in_features), + .alpha = 1.0f, + .beta = 0.0f, + .batch_count = 1, + .input_dtype = compute_dtype, + .output_dtype = output_dtype, + }); + }), + DataType::kBFLOAT16) + default: + LOG_UNSUPPORTED_DTYPE(compute_dtype, "MACA LinearBackwardInput"); + } + + 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) { + const auto &grad_output_dims = grad_output->Dims(); + CHECK_GE(grad_output_dims.size(), 2); + const int64_t bs = std::accumulate(grad_output_dims.rbegin() + 1, grad_output_dims.rend(), int64_t{1}, + std::multiplies{}); + + 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)); + + // FIXME(cx): output dtype promotion is a temporary hack; revisit when autograd/autocast is fixed. + auto output_dtype = (compute_dtype == DataType::kBFLOAT16) ? DataType::kFLOAT32 : compute_dtype; + const std::vector weight_dims + = transpose ? std::vector{out_features, in_features} : std::vector{in_features, out_features}; + auto grad_weight = std::make_shared(weight_dims, output_dtype, grad_output->GetDevice()); + + auto device = grad_output->GetDevice(); + const int m = transpose ? in_features : out_features; + const int n = transpose ? out_features : in_features; + const void *a = transpose ? input->DataPtr() : grad_output_promoted->DataPtr(); + const void *b = transpose ? grad_output_promoted->DataPtr() : input->DataPtr(); + const int lda = transpose ? in_features : out_features; + const int ldb = transpose ? out_features : in_features; + const int ldc = transpose ? in_features : out_features; + + Dispatcher::Instance().Call({device.type(), "Gemm"}, device, + GemmParams{ + .trans_a = GemmTranspose::kNoTranspose, + .trans_b = GemmTranspose::kTranspose, + .m = m, + .n = n, + .k = static_cast(bs), + .A = a, + .lda = lda, + .B = b, + .ldb = ldb, + .C = grad_weight->DataPtr(), + .ldc = ldc, + .alpha = 1.0f, + .beta = 0.0f, + .batch_count = 1, + .input_dtype = compute_dtype, + .output_dtype = output_dtype, + }); + + return grad_weight; +} + +std::shared_ptr LinearBackwardBias(const std::shared_ptr &grad_output, int64_t out_features) { + const auto &dims = grad_output->Dims(); + CHECK_GE(dims.size(), 2); + CHECK_EQ(dims.back(), out_features); + const int64_t bs = std::accumulate(dims.rbegin() + 1, dims.rend(), int64_t{1}, std::multiplies{}); + + auto compute_dtype = grad_output->Dtype(); + // FIXME(cx): output dtype promotion is a temporary hack; revisit when autograd/autocast is fixed. + auto output_dtype = (compute_dtype == DataType::kBFLOAT16) ? DataType::kFLOAT32 : compute_dtype; + auto grad_bias + = std::make_shared(std::vector{out_features}, output_dtype, grad_output->GetDevice()); + + auto device = grad_output->GetDevice(); + const auto &maca_stream = dynamic_cast( + infini_train::core::GetDeviceGuardImpl(device.type())->GetStream(device)) + ->maca_stream(); + + constexpr int BLOCK_SIZE = 256; + switch (compute_dtype) { + DISPATCH_CASE(WRAP({ + ReduceRowsKernel<<>>( + static_cast(grad_output->DataPtr()), + static_cast(grad_bias->DataPtr()), bs, out_features); + }), + DataType::kFLOAT32) + DISPATCH_CASE(WRAP({ + ReduceRowsKernel<<>>( + static_cast(grad_output->DataPtr()), + static_cast(grad_bias->DataPtr()), bs, out_features); + }), + DataType::kBFLOAT16) + default: + LOG_UNSUPPORTED_DTYPE(compute_dtype, "MACA LinearBackwardBias"); + } + + return grad_bias; +} +} // namespace infini_train::kernels::maca + +namespace infini_train::kernels::maca { + +void RegisterLinearKernels() { + REGISTER_KERNEL(infini_train::Device::DeviceType::kPrivateUse1, MatmulForward, + infini_train::kernels::maca::MatmulForward) + REGISTER_KERNEL(infini_train::Device::DeviceType::kPrivateUse1, MatmulBackwardInput, + infini_train::kernels::maca::MatmulBackwardInput) + REGISTER_KERNEL(infini_train::Device::DeviceType::kPrivateUse1, MatmulBackwardOther, + infini_train::kernels::maca::MatmulBackwardOther) + REGISTER_KERNEL(infini_train::Device::DeviceType::kPrivateUse1, LinearForward, + infini_train::kernels::maca::LinearForward) + REGISTER_KERNEL(infini_train::Device::DeviceType::kPrivateUse1, LinearBackwardInput, + infini_train::kernels::maca::LinearBackwardInput) + REGISTER_KERNEL(infini_train::Device::DeviceType::kPrivateUse1, LinearBackwardWeight, + infini_train::kernels::maca::LinearBackwardWeight) + REGISTER_KERNEL(infini_train::Device::DeviceType::kPrivateUse1, LinearBackwardBias, + infini_train::kernels::maca::LinearBackwardBias) +} + +} // namespace infini_train::kernels::maca diff --git a/backends/maca/src/kernels/no_op.maca b/backends/maca/src/kernels/no_op.maca new file mode 100644 index 0000000..fe66df6 --- /dev/null +++ b/backends/maca/src/kernels/no_op.maca @@ -0,0 +1,39 @@ +#include +#include +#include +#include +#include + +#include "glog/logging.h" + +#include "infini_train/include/dispatcher.h" +#include "infini_train/include/tensor.h" + +namespace infini_train::kernels::maca { +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::maca + +namespace infini_train::kernels::maca { + +void RegisterNoOpKernels() { + REGISTER_KERNEL(infini_train::Device::DeviceType::kPrivateUse1, NoOpForward, + infini_train::kernels::maca::NoOpForward) + REGISTER_KERNEL(infini_train::Device::DeviceType::kPrivateUse1, NoOpBackward, + infini_train::kernels::maca::NoOpBackward) +} + +} // namespace infini_train::kernels::maca diff --git a/backends/maca/src/kernels/outer.maca b/backends/maca/src/kernels/outer.maca new file mode 100644 index 0000000..aa6e8cb --- /dev/null +++ b/backends/maca/src/kernels/outer.maca @@ -0,0 +1,195 @@ +#include +#include +#include +#include + +#include + +#include "glog/logging.h" + +#include "common/common_maca.h" +#include "infini_train/include/core/runtime/device_guard.h" +#include "infini_train/include/dispatcher.h" +#include "infini_train/include/tensor.h" + +#include "kernels/common/gemm.h" +#include "runtime/maca_dispatch.h" +#include "runtime/maca_runtime_common.h" + +namespace infini_train::kernels::maca { +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] + Dispatcher::Instance().Call({device.type(), "Gemm"}, device, + GemmParams{ + .trans_a = GemmTranspose::kNoTranspose, + .trans_b = GemmTranspose::kNoTranspose, + .m = static_cast(N), + .n = static_cast(M), + .k = 1, + .A = other->DataPtr(), + .lda = static_cast(N), + .B = input->DataPtr(), + .ldb = 1, + .C = output->DataPtr(), + .ldc = static_cast(N), + .alpha = 1.0f, + .beta = 0.0f, + .batch_count = 1, + .input_dtype = input->Dtype(), + .output_dtype = output->Dtype(), + }); + + 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()); + + auto device = input->GetDevice(); + float alpha = 1.0f; + float beta = 0.0f; + mcblasHandle_t handle = dynamic_cast( + infini_train::core::GetDeviceGuardImpl(device.type())->GetBlasHandle(device)) + ->mcblas_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] + MCBLAS_CHECK(mcblasSgemv(handle, MCBLAS_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] + MCBLAS_CHECK(mcblasSgemv(handle, MCBLAS_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( + // mcblasSgemv does not support bf16; use the generic GEMM kernel instead. + 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] + Dispatcher::Instance().Call({device.type(), "Gemm"}, device, + GemmParams{ + .trans_a = GemmTranspose::kNoTranspose, + .trans_b = GemmTranspose::kNoTranspose, + .m = 1, + .n = static_cast(M), + .k = static_cast(N), + .A = other_promoted->DataPtr(), + .lda = 1, + .B = grad_output_promoted->DataPtr(), + .ldb = static_cast(N), + .C = grad_input->DataPtr(), + .ldc = 1, + .alpha = 1.0f, + .beta = 0.0f, + .batch_count = 1, + .input_dtype = promoted_type, + .output_dtype = output_dtype, + }); + // 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] + Dispatcher::Instance().Call({device.type(), "Gemm"}, device, + GemmParams{ + .trans_a = GemmTranspose::kNoTranspose, + .trans_b = GemmTranspose::kTranspose, + .m = 1, + .n = static_cast(N), + .k = static_cast(M), + .A = input_promoted->DataPtr(), + .lda = 1, + .B = grad_output_promoted->DataPtr(), + .ldb = static_cast(N), + .C = grad_other->DataPtr(), + .ldc = 1, + .alpha = 1.0f, + .beta = 0.0f, + .batch_count = 1, + .input_dtype = promoted_type, + .output_dtype = output_dtype, + }); + }), + DataType::kBFLOAT16) + default: + LOG_UNSUPPORTED_DTYPE(promoted_type, "MACA OuterBackward"); + } + + return {grad_input, grad_other}; +} + +} // namespace infini_train::kernels::maca + +namespace infini_train::kernels::maca { + +void RegisterOuterKernels() { + REGISTER_KERNEL(infini_train::Device::DeviceType::kPrivateUse1, OuterForward, + infini_train::kernels::maca::OuterForward) + REGISTER_KERNEL(infini_train::Device::DeviceType::kPrivateUse1, OuterBackward, + infini_train::kernels::maca::OuterBackward) +} + +} // namespace infini_train::kernels::maca diff --git a/backends/maca/src/kernels/reduction.maca b/backends/maca/src/kernels/reduction.maca new file mode 100644 index 0000000..3a0157a --- /dev/null +++ b/backends/maca/src/kernels/reduction.maca @@ -0,0 +1,258 @@ +#include +#include +#include +#include +#include +#include +#include + +#include + +#include "infini_train/include/core/runtime/device_guard.h" +#include "infini_train/include/dispatcher.h" +#include "infini_train/include/tensor.h" + +#include "common/common_maca.h" +#include "common/cub_compat.cuh" +#include "common/kernel_helper.cuh" +#include "runtime/maca_dispatch.h" +#include "runtime/maca_runtime_common.h" + +namespace infini_train::kernels::maca { +namespace { +constexpr float kInfinity = std::numeric_limits::infinity(); +} // namespace + +namespace { +// Reduction operators +template struct CubOp; + +template struct CubOp { + __device__ static T Init() { return common::maca::Cast(0); } + __device__ static T Reduce(T a, T b) { return common::maca::Add(a, b); } + __device__ static CubSumOp Op() { return CubSumOp(); } +}; + +template struct CubOp { + __device__ static T Init() { return common::maca::Cast(-kInfinity); } + __device__ static T Reduce(T a, T b) { return common::maca::Max(a, b); } + __device__ static CubMaxOp Op() { return CubMaxOp(); } +}; + +template struct CubOp { + __device__ static T Init() { return common::maca::Cast(kInfinity); } + __device__ static T Reduce(T a, T b) { return common::maca::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::maca::Div(sum, common::maca::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 = cub::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 __maca_bfloat16) <-> (integral types) is supported + if constexpr (std::is_same_v || std::is_same_v) { + H_casted = common::maca::Cast(static_cast(H)); + } else { + H_casted = common::maca::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 &maca_stream = dynamic_cast( + infini_train::core::GetDeviceGuardImpl(device.type())->GetStream(device)) + ->maca_stream(); + + core::maca::DispatchMacaFunc( + dtype, + [=]() { + GenericReduceKernel, BLOCK_SIZE> + <<>>(static_cast(input->DataPtr()), + static_cast(output->DataPtr()), N, H, W, + FinalizeOp{}); + }, + "MACA 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 &maca_stream = dynamic_cast( + infini_train::core::GetDeviceGuardImpl(device.type())->GetStream(device)) + ->maca_stream(); + + core::maca::DispatchMacaFunc( + dtype, + [=]() { + // The backward kernel assigns every grad_input element on all reduction branches; no Fill is needed. + 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); + }, + "MACA 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::maca + +namespace infini_train::kernels::maca { + +void RegisterReductionKernels() { + REGISTER_KERNEL(infini_train::Device::DeviceType::kPrivateUse1, MeanForward, + infini_train::kernels::maca::MeanForward) + REGISTER_KERNEL(infini_train::Device::DeviceType::kPrivateUse1, SumForward, infini_train::kernels::maca::SumForward) + REGISTER_KERNEL(infini_train::Device::DeviceType::kPrivateUse1, MaxForward, infini_train::kernels::maca::MaxForward) + REGISTER_KERNEL(infini_train::Device::DeviceType::kPrivateUse1, MinForward, infini_train::kernels::maca::MinForward) + REGISTER_KERNEL(infini_train::Device::DeviceType::kPrivateUse1, MeanBackward, + infini_train::kernels::maca::MeanBackward) + REGISTER_KERNEL(infini_train::Device::DeviceType::kPrivateUse1, SumBackward, + infini_train::kernels::maca::SumBackward) + REGISTER_KERNEL(infini_train::Device::DeviceType::kPrivateUse1, MaxBackward, + infini_train::kernels::maca::MaxBackward) + REGISTER_KERNEL(infini_train::Device::DeviceType::kPrivateUse1, MinBackward, + infini_train::kernels::maca::MinBackward) +} + +} // namespace infini_train::kernels::maca diff --git a/backends/maca/src/kernels/register_maca_kernels.cc b/backends/maca/src/kernels/register_maca_kernels.cc new file mode 100644 index 0000000..0b71a19 --- /dev/null +++ b/backends/maca/src/kernels/register_maca_kernels.cc @@ -0,0 +1,53 @@ +#include "kernels/register_maca_kernels.h" + +namespace infini_train::kernels::maca { + +void RegisterAccumulateGradKernels(); +void RegisterCastKernels(); +void RegisterConcatKernels(); +void RegisterCrossEntropyKernels(); +void RegisterElementwiseKernels(); +void RegisterEmbeddingKernels(); +void RegisterFillKernels(); +void RegisterGatherKernels(); +void RegisterGemmKernels(); +void RegisterLayerNormKernels(); +void RegisterLinearKernels(); +void RegisterNoOpKernels(); +void RegisterOuterKernels(); +void RegisterReductionKernels(); +void RegisterScatterKernels(); +void RegisterSliceKernels(); +void RegisterSoftmaxKernels(); +void RegisterSplitKernels(); +void RegisterStackKernels(); +void RegisterTopKKernels(); +void RegisterTransformKernels(); +void RegisterVocabParallelCrossEntropyKernels(); + +void RegisterMacaKernels() { + RegisterAccumulateGradKernels(); + RegisterCastKernels(); + RegisterConcatKernels(); + RegisterCrossEntropyKernels(); + RegisterElementwiseKernels(); + RegisterEmbeddingKernels(); + RegisterFillKernels(); + RegisterGatherKernels(); + RegisterGemmKernels(); + RegisterLayerNormKernels(); + RegisterLinearKernels(); + RegisterNoOpKernels(); + RegisterOuterKernels(); + RegisterReductionKernels(); + RegisterScatterKernels(); + RegisterSliceKernels(); + RegisterSoftmaxKernels(); + RegisterSplitKernels(); + RegisterStackKernels(); + RegisterTopKKernels(); + RegisterTransformKernels(); + RegisterVocabParallelCrossEntropyKernels(); +} + +} // namespace infini_train::kernels::maca diff --git a/backends/maca/src/kernels/register_maca_kernels.h b/backends/maca/src/kernels/register_maca_kernels.h new file mode 100644 index 0000000..611b110 --- /dev/null +++ b/backends/maca/src/kernels/register_maca_kernels.h @@ -0,0 +1,7 @@ +#pragma once + +namespace infini_train::kernels::maca { + +void RegisterMacaKernels(); + +} // namespace infini_train::kernels::maca diff --git a/backends/maca/src/kernels/scatter.maca b/backends/maca/src/kernels/scatter.maca new file mode 100644 index 0000000..073cb42 --- /dev/null +++ b/backends/maca/src/kernels/scatter.maca @@ -0,0 +1,126 @@ +#include +#include +#include +#include +#include + +#include "glog/logging.h" + +#include "common/common_maca.h" +#include "infini_train/include/core/runtime/device_guard.h" +#include "infini_train/include/datatype.h" +#include "infini_train/include/dispatcher.h" +#include "infini_train/include/tensor.h" + +#include "runtime/maca_dispatch.h" +#include "runtime/maca_runtime_common.h" + +namespace infini_train::kernels::maca { + +template +__global__ void ScatterForwardKernel(const T *__restrict__ values, const int64_t *__restrict__ indices, + T *__restrict__ output, int64_t rows, int64_t topk, int64_t num_experts) { + int64_t idx = blockIdx.x * blockDim.x + threadIdx.x; + const int64_t total = rows * topk; + if (idx >= total) { + return; + } + + const int64_t row = idx / topk; + const int64_t expert_idx = indices[idx]; + output[row * num_experts + expert_idx] = values[idx]; +} + +std::shared_ptr ScatterForward(const std::shared_ptr &values, const std::shared_ptr &indices, + const std::vector &output_dims) { + CHECK(indices->Dtype() == DataType::kINT64) << "MACA ScatterForward expects int64 indices"; + CHECK(values->Dims() == indices->Dims()); + CHECK(!output_dims.empty()); + CHECK_EQ(values->Dims().size(), output_dims.size()); + CHECK_GT(values->Dims().back(), 0); + CHECK_GT(output_dims.back(), 0); + + const int64_t topk = values->Dims().back(); + const int64_t num_experts = output_dims.back(); + const int64_t rows = static_cast(values->NumElements()) / topk; + const int64_t output_numel + = std::accumulate(output_dims.begin(), output_dims.end(), int64_t{1}, std::multiplies()); + CHECK_EQ(output_numel, rows * num_experts); + + auto output = std::make_shared(output_dims, values->Dtype(), values->GetDevice()); + + auto device = values->GetDevice(); + const auto &stream = dynamic_cast( + infini_train::core::GetDeviceGuardImpl(device.type())->GetStream(device)) + ->maca_stream(); + MACA_CHECK(mcMemsetAsync(output->DataPtr(), 0, output->SizeInBytes(), stream)); + constexpr int threads = 256; + const int blocks = static_cast(((rows * topk) + threads - 1) / threads); + core::maca::DispatchMacaFunc>( + {values->Dtype()}, + [=]() { + ScatterForwardKernel<<>>( + static_cast(values->DataPtr()), static_cast(indices->DataPtr()), + static_cast(output->DataPtr()), rows, topk, num_experts); + MACA_CHECK(mcGetLastError()); + }, + "MACA ScatterForward"); + return output; +} + +template +__global__ void ScatterBackwardKernel(const T *__restrict__ grad_output, const int64_t *__restrict__ indices, + T *__restrict__ grad_values, int64_t rows, int64_t topk, int64_t num_experts) { + int64_t idx = blockIdx.x * blockDim.x + threadIdx.x; + const int64_t total = rows * topk; + if (idx >= total) { + return; + } + const int64_t row = idx / topk; + const int64_t expert_idx = indices[idx]; + grad_values[idx] = grad_output[row * num_experts + expert_idx]; +} + +std::shared_ptr ScatterBackward(const std::shared_ptr &grad_output, + const std::shared_ptr &indices) { + CHECK(indices->Dtype() == DataType::kINT64) << "MACA ScatterBackward expects int64 indices"; + CHECK(IsFloatingPointDType(grad_output->Dtype())) + << "MACA ScatterBackward only supports floating grad_output dtype, got " + << kDataTypeToDesc.at(grad_output->Dtype()); + CHECK_GE(grad_output->Dims().size(), 1); + CHECK_GE(indices->Dims().size(), 1); + const int64_t num_experts = grad_output->Dims().back(); + const int64_t topk = indices->Dims().back(); + const int64_t rows = static_cast(indices->NumElements()) / topk; + CHECK_EQ(grad_output->NumElements(), static_cast(rows * num_experts)); + + auto grad_values = std::make_shared(indices->Dims(), grad_output->Dtype(), grad_output->GetDevice()); + + auto device = grad_output->GetDevice(); + const auto &stream = dynamic_cast( + infini_train::core::GetDeviceGuardImpl(device.type())->GetStream(device)) + ->maca_stream(); + constexpr int threads = 256; + const int blocks = static_cast(((rows * topk) + threads - 1) / threads); + + core::maca::DispatchMacaFunc( + grad_output->Dtype(), + [=]() { + ScatterBackwardKernel<<>>( + static_cast(grad_output->DataPtr()), static_cast(indices->DataPtr()), + static_cast(grad_values->DataPtr()), rows, topk, num_experts); + MACA_CHECK(mcGetLastError()); + }, + "MACA ScatterBackward"); + + return grad_values; +} + +void RegisterScatterKernels() { + REGISTER_KERNEL(infini_train::Device::DeviceType::kPrivateUse1, ScatterForward, + infini_train::kernels::maca::ScatterForward) + REGISTER_KERNEL(infini_train::Device::DeviceType::kPrivateUse1, ScatterBackward, + infini_train::kernels::maca::ScatterBackward) +} + +} // namespace infini_train::kernels::maca diff --git a/backends/maca/src/kernels/slice.maca b/backends/maca/src/kernels/slice.maca new file mode 100644 index 0000000..8d92185 --- /dev/null +++ b/backends/maca/src/kernels/slice.maca @@ -0,0 +1,211 @@ +#include +#include +#include + +#include "glog/logging.h" + +#include "common/common_maca.h" +#include "infini_train/include/core/runtime/device_guard.h" +#include "infini_train/include/dispatcher.h" +#include "infini_train/include/tensor.h" + +#include "runtime/maca_dispatch.h" +#include "runtime/maca_runtime_common.h" + +namespace infini_train::kernels::maca { + +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()); + // SliceForwardKernel writes every output index in [0, total_elements); no Fill is needed. + + 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)) + ->maca_stream(); + + mcMallocAsync(reinterpret_cast(&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(); + + mcMemcpyAsync(new_dims_dev, new_dims.data(), ends.size() * sizeof(int64_t), mcMemcpyHostToDevice, stream); + mcMemcpyAsync(starts_dev, starts.data(), starts.size() * sizeof(int64_t), mcMemcpyHostToDevice, stream); + mcMemcpyAsync(steps_dev, steps.data(), steps.size() * sizeof(int64_t), mcMemcpyHostToDevice, stream); + mcMemcpyAsync(input_strides_dev, src_strides.data(), dims.size() * sizeof(int64_t), mcMemcpyHostToDevice, stream); + mcMemcpyAsync(output_strides_dev, dst_strides.data(), new_dims.size() * sizeof(int64_t), mcMemcpyHostToDevice, + stream); + + int threads_per_block = 256; + int num_blocks = (total_elements + threads_per_block - 1) / threads_per_block; + + core::maca::DispatchMacaFunc( + 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); + }, + "MACA SliceForward"); + + mcFreeAsync(new_dims_dev, stream); + + 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()); + core::maca::DispatchMacaFunc( + grad_output_dtype, [=]() { grad_input->Fill(0); }, "MACA SliceBackward"); + + 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)) + ->maca_stream(); + mcMallocAsync(reinterpret_cast(&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(); + + mcMemcpyAsync(new_dims_dev, new_dims.data(), ends.size() * sizeof(int64_t), mcMemcpyHostToDevice, stream); + mcMemcpyAsync(starts_dev, starts.data(), starts.size() * sizeof(int64_t), mcMemcpyHostToDevice, stream); + mcMemcpyAsync(steps_dev, steps.data(), steps.size() * sizeof(int64_t), mcMemcpyHostToDevice, stream); + mcMemcpyAsync(input_strides_dev, src_strides.data(), dims.size() * sizeof(int64_t), mcMemcpyHostToDevice, stream); + mcMemcpyAsync(output_strides_dev, dst_strides.data(), new_dims.size() * sizeof(int64_t), mcMemcpyHostToDevice, + stream); + + int threads_per_block = 256; + int num_blocks = (total_elements + threads_per_block - 1) / threads_per_block; + + core::maca::DispatchMacaFunc( + 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); + }, + "MACA SliceBackward"); + + mcFreeAsync(new_dims_dev, stream); + + return grad_input; +} +} // namespace infini_train::kernels::maca + +namespace infini_train::kernels::maca { + +void RegisterSliceKernels() { + REGISTER_KERNEL(infini_train::Device::DeviceType::kPrivateUse1, SliceForward, + infini_train::kernels::maca::SliceForward) + REGISTER_KERNEL(infini_train::Device::DeviceType::kPrivateUse1, SliceBackward, + infini_train::kernels::maca::SliceBackward) +} + +} // namespace infini_train::kernels::maca diff --git a/backends/maca/src/kernels/softmax.maca b/backends/maca/src/kernels/softmax.maca new file mode 100644 index 0000000..2703491 --- /dev/null +++ b/backends/maca/src/kernels/softmax.maca @@ -0,0 +1,228 @@ +#include +#include +#include +#include + +#include + +#include "glog/logging.h" + +#include "common/common_maca.h" +#include "common/cub_compat.cuh" +#include "common/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 "runtime/maca_dispatch.h" +#include "runtime/maca_runtime_common.h" + +namespace infini_train::kernels::maca { +template +__global__ void SoftmaxForwardKernel(T *output, const T *input, int64_t outer_size, int64_t axis_size, + int64_t inner_size) { + using BlockReduce = cub::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::maca::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::maca::Cast(input[idx]) - row_max); + output[idx] = common::maca::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::maca::Cast(common::maca::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, "MACA 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, "MACA 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 &maca_stream = dynamic_cast( + infini_train::core::GetDeviceGuardImpl(device.type())->GetStream(device)) + ->maca_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, __maca_bfloat16>(output, input, dim);), DataType::kBFLOAT16) + default: + LOG_LOC(FATAL, "MACA 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 = cub::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::maca::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::maca::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, "MACA 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, "MACA 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 &maca_stream = dynamic_cast( + infini_train::core::GetDeviceGuardImpl(device.type())->GetStream(device)) + ->maca_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()); + // For non-empty tensors, the grid covers every outer/axis/inner index exactly once; no Fill is needed. + + switch (promoted_type) { + DISPATCH_CASE(WRAP(LaunchBackward<256, float>(grad_input, grad_output_promoted, output_promoted, dim);), + DataType::kFLOAT32) + DISPATCH_CASE( + WRAP(LaunchBackward<256, __maca_bfloat16>(grad_input, grad_output_promoted, output_promoted, dim);), + DataType::kBFLOAT16) + default: + LOG_LOC(FATAL, "MACA softmax backward: 'Unsupported data type'"); + } + + return grad_input; +} +} // namespace infini_train::kernels::maca + +namespace infini_train::kernels::maca { + +void RegisterSoftmaxKernels() { + REGISTER_KERNEL(infini_train::Device::DeviceType::kPrivateUse1, SoftmaxForward, + infini_train::kernels::maca::SoftmaxForward) + REGISTER_KERNEL(infini_train::Device::DeviceType::kPrivateUse1, SoftmaxBackward, + infini_train::kernels::maca::SoftmaxBackward) +} + +} // namespace infini_train::kernels::maca diff --git a/backends/maca/src/kernels/split.maca b/backends/maca/src/kernels/split.maca new file mode 100644 index 0000000..c7370e1 --- /dev/null +++ b/backends/maca/src/kernels/split.maca @@ -0,0 +1,190 @@ +#include +#include +#include +#include +#include +#include +#include + +#include "common/common_maca.h" +#include "infini_train/include/core/runtime/device_guard.h" +#include "infini_train/include/dispatcher.h" +#include "infini_train/include/tensor.h" + +#include "runtime/maca_dispatch.h" +#include "runtime/maca_runtime_common.h" + +namespace infini_train::kernels::maca { +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 &maca_stream = dynamic_cast( + infini_train::core::GetDeviceGuardImpl(device.type())->GetStream(device)) + ->maca_stream(); + + core::maca::DispatchMacaFunc( + dtype, + [=]() { + SplitForwardKernel<<>>( + static_cast(input->DataPtr()), static_cast(output->DataPtr()), N, H_in, H_out, W, + start); + }, + "MACA 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()); + // Keep initialization: defensive early returns in SplitBackwardKernel can leave elements unwritten. + grad_input->Fill(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)) + ->maca_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; + mcMallocAsync(reinterpret_cast(&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); + + mcMemcpyAsync(device_grad_output_ptrs, host_grad_output_ptrs.data(), sizeof(T *) * num_splits, mcMemcpyHostToDevice, + 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); } + + mcMemcpyAsync(device_H_outs, H_outs.data(), sizeof(int64_t) * num_splits, mcMemcpyHostToDevice, 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); + + mcFreeAsync(device_ptr, stream); + + 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::maca::DispatchMacaFunc( + grad_outputs[0]->Dtype(), + [=]() { return LaunchSplitBackward(input_dims, split_size, dim, grad_outputs); }, + "MACA SplitBackward"); +} +} // namespace infini_train::kernels::maca + +namespace infini_train::kernels::maca { + +void RegisterSplitKernels() { + REGISTER_KERNEL(infini_train::Device::DeviceType::kPrivateUse1, SplitForward, + infini_train::kernels::maca::SplitForward) + REGISTER_KERNEL(infini_train::Device::DeviceType::kPrivateUse1, SplitBackward, + infini_train::kernels::maca::SplitBackward) +} + +} // namespace infini_train::kernels::maca diff --git a/backends/maca/src/kernels/stack.maca b/backends/maca/src/kernels/stack.maca new file mode 100644 index 0000000..f9b0c4c --- /dev/null +++ b/backends/maca/src/kernels/stack.maca @@ -0,0 +1,165 @@ +#include +#include +#include +#include +#include +#include +#include + +#include "glog/logging.h" + +#include "common/common_maca.h" +#include "infini_train/include/core/runtime/device_guard.h" +#include "infini_train/include/dispatcher.h" +#include "infini_train/include/tensor.h" + +#include "runtime/maca_dispatch.h" +#include "runtime/maca_runtime_common.h" + +namespace infini_train::kernels::maca { +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)) + ->maca_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::maca::DispatchMacaFunc( + 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; + mcMallocAsync(reinterpret_cast(&device_input_ptrs), sizeof(T *) * num_inputs, stream); + mcMemcpyAsync(device_input_ptrs, host_input_ptrs.data(), sizeof(T *) * num_inputs, mcMemcpyHostToDevice, + stream); + + StackForwardKernel<<>>( + device_input_ptrs, static_cast(output->DataPtr()), N, D, num_inputs); + + mcFreeAsync(device_input_ptrs, stream); + }, + "MACA 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()); + // StackBackwardKernel writes every element of every grad tensor exactly once; no Fill is needed. + 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)) + ->maca_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::maca::DispatchMacaFunc( + dtype, + [=]() { + std::vector host_ptrs; + for (auto &t : grads) { host_ptrs.push_back(static_cast(t->DataPtr())); } + + T **device_ptrs; + mcMallocAsync(reinterpret_cast(&device_ptrs), sizeof(T *) * num_inputs, stream); + mcMemcpyAsync(device_ptrs, host_ptrs.data(), sizeof(T *) * num_inputs, mcMemcpyHostToDevice, stream); + + StackBackwardKernel<<>>( + static_cast(grad_output->DataPtr()), device_ptrs, N, D, num_inputs); + + mcFreeAsync(device_ptrs, stream); + }, + "MACA StackBackward"); + + return grads; +} + +} // namespace infini_train::kernels::maca + +namespace infini_train::kernels::maca { + +void RegisterStackKernels() { + REGISTER_KERNEL(infini_train::Device::DeviceType::kPrivateUse1, StackForward, + infini_train::kernels::maca::StackForward) + REGISTER_KERNEL(infini_train::Device::DeviceType::kPrivateUse1, StackBackward, + infini_train::kernels::maca::StackBackward) +} + +} // namespace infini_train::kernels::maca diff --git a/backends/maca/src/kernels/topk.maca b/backends/maca/src/kernels/topk.maca new file mode 100644 index 0000000..f4313fa --- /dev/null +++ b/backends/maca/src/kernels/topk.maca @@ -0,0 +1,161 @@ +#include +#include +#include + +#include "glog/logging.h" + +#include "common/common_maca.h" +#include "common/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 "runtime/maca_dispatch.h" +#include "runtime/maca_runtime_common.h" + +namespace infini_train::kernels::maca { + +template +__global__ void TopKForwardKernel(const T *__restrict__ input, T *__restrict__ top_values, + int64_t *__restrict__ top_indices, int64_t rows, int64_t dim_size, int64_t inner_size, + int64_t topk, bool largest) { + int64_t row = blockIdx.x * blockDim.x + threadIdx.x; + if (row >= rows) { + return; + } + + const int64_t outer = row / inner_size; + const int64_t inner = row % inner_size; + for (int64_t idx = 0; idx < dim_size; ++idx) { + const float value = common::maca::Cast(input[outer * dim_size * inner_size + idx * inner_size + inner]); + int64_t rank = 0; + for (int64_t other_idx = 0; other_idx < dim_size; ++other_idx) { + const float other_value + = common::maca::Cast(input[outer * dim_size * inner_size + other_idx * inner_size + inner]); + const bool ranks_before = largest ? (other_value > value || (other_value == value && other_idx < idx)) + : (other_value < value || (other_value == value && other_idx < idx)); + if (ranks_before) { + ++rank; + } + } + if (rank < topk) { + const int64_t out_offset = outer * topk * inner_size + rank * inner_size + inner; + top_values[out_offset] = input[outer * dim_size * inner_size + idx * inner_size + inner]; + top_indices[out_offset] = idx; + } + } +} + +std::vector> TopKForward(const std::shared_ptr &input, int64_t topk, int64_t dim, + bool largest, bool sorted) { + CHECK_GE(input->Dims().size(), 1); + CHECK(sorted) << "TopK currently only supports sorted=true"; + const auto &dims = input->Dims(); + if (dim < 0) { + dim += static_cast(dims.size()); + } + CHECK_GE(dim, 0); + CHECK_LT(dim, static_cast(dims.size())); + + const int64_t dim_size = dims[dim]; + CHECK_GT(dim_size, 0); + CHECK_GT(topk, 0); + CHECK_LE(topk, dim_size); + int64_t outer_size = 1; + for (int64_t idx = 0; idx < dim; ++idx) { outer_size *= dims[idx]; } + int64_t inner_size = 1; + for (size_t idx = static_cast(dim) + 1; idx < dims.size(); ++idx) { inner_size *= dims[idx]; } + const int64_t rows = outer_size * inner_size; + + auto topk_dims = dims; + topk_dims[dim] = topk; + auto top_values = std::make_shared(topk_dims, input->Dtype(), input->GetDevice()); + auto top_indices = std::make_shared(topk_dims, DataType::kINT64, input->GetDevice()); + + auto device = input->GetDevice(); + const auto &stream = dynamic_cast( + infini_train::core::GetDeviceGuardImpl(device.type())->GetStream(device)) + ->maca_stream(); + constexpr int threads = 256; + const int blocks = static_cast((rows + threads - 1) / threads); + + core::maca::DispatchMacaFunc( + input->Dtype(), + [=]() { + TopKForwardKernel<<>>( + static_cast(input->DataPtr()), static_cast(top_values->DataPtr()), + static_cast(top_indices->DataPtr()), rows, dim_size, inner_size, topk, largest); + MACA_CHECK(mcGetLastError()); + }, + "MACA TopKForward"); + + return {top_values, top_indices}; +} + +template +__global__ void TopKBackwardKernel(const T *__restrict__ grad_values, const int64_t *__restrict__ indices, + T *__restrict__ grad_input, int64_t rows, int64_t dim_size, int64_t inner_size, + int64_t topk) { + int64_t row = blockIdx.x * blockDim.x + threadIdx.x; + if (row >= rows) { + return; + } + + const int64_t outer = row / inner_size; + const int64_t inner = row % inner_size; + for (int64_t selected = 0; selected < topk; ++selected) { + const int64_t out_offset = outer * topk * inner_size + selected * inner_size + inner; + const int64_t selected_idx = indices[out_offset]; + grad_input[outer * dim_size * inner_size + selected_idx * inner_size + inner] = grad_values[out_offset]; + } +} + +std::shared_ptr TopKBackward(const std::shared_ptr &grad_values, const std::shared_ptr &indices, + const std::vector &input_dims, int64_t dim) { + CHECK(indices->Dtype() == DataType::kINT64) << "MACA TopKBackward expects int64 indices"; + CHECK(grad_values->Dims() == indices->Dims()); + CHECK(!input_dims.empty()); + if (dim < 0) { + dim += static_cast(input_dims.size()); + } + CHECK_GE(dim, 0); + CHECK_LT(dim, static_cast(input_dims.size())); + + const int64_t dim_size = input_dims[dim]; + const int64_t topk = indices->Dims()[dim]; + int64_t outer_size = 1; + for (int64_t idx = 0; idx < dim; ++idx) { outer_size *= input_dims[idx]; } + int64_t inner_size = 1; + for (size_t idx = static_cast(dim) + 1; idx < input_dims.size(); ++idx) { inner_size *= input_dims[idx]; } + const int64_t rows = outer_size * inner_size; + + auto grad_input = std::make_shared(input_dims, grad_values->Dtype(), grad_values->GetDevice()); + auto device = grad_values->GetDevice(); + const auto &stream = dynamic_cast( + infini_train::core::GetDeviceGuardImpl(device.type())->GetStream(device)) + ->maca_stream(); + MACA_CHECK(mcMemsetAsync(grad_input->DataPtr(), 0, grad_input->SizeInBytes(), stream)); + + constexpr int threads = 256; + const int blocks = static_cast((rows + threads - 1) / threads); + core::maca::DispatchMacaFunc( + grad_values->Dtype(), + [=]() { + TopKBackwardKernel<<>>( + static_cast(grad_values->DataPtr()), static_cast(indices->DataPtr()), + static_cast(grad_input->DataPtr()), rows, dim_size, inner_size, topk); + MACA_CHECK(mcGetLastError()); + }, + "MACA TopKBackward"); + + return grad_input; +} + +void RegisterTopKKernels() { + REGISTER_KERNEL(infini_train::Device::DeviceType::kPrivateUse1, TopKForward, + infini_train::kernels::maca::TopKForward) + REGISTER_KERNEL(infini_train::Device::DeviceType::kPrivateUse1, TopKBackward, + infini_train::kernels::maca::TopKBackward) +} + +} // namespace infini_train::kernels::maca diff --git a/backends/maca/src/kernels/transform.maca b/backends/maca/src/kernels/transform.maca new file mode 100644 index 0000000..c0c5a5b --- /dev/null +++ b/backends/maca/src/kernels/transform.maca @@ -0,0 +1,607 @@ +#include +#include +#include +#include +#include +#include +#include +#include + +#include "common/common_maca.h" +#include "common/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 "runtime/maca_dispatch.h" +#include "runtime/maca_runtime_common.h" + +namespace infini_train::kernels::maca { + +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 &maca_stream = dynamic_cast( + infini_train::core::GetDeviceGuardImpl(device.type())->GetStream(device)) + ->maca_stream(); + + core::maca::DispatchMacaFunc( + input->Dtype(), + [=]() { + TrilForwardKernel<<>>( + static_cast(input->DataPtr()), static_cast(output->DataPtr()), rows, cols, diagonal); + }, + "MACA 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 &maca_stream = dynamic_cast( + infini_train::core::GetDeviceGuardImpl(device.type())->GetStream(device)) + ->maca_stream(); + + core::maca::DispatchMacaFunc( + dtype, + [=]() { + grad_input->Fill(0); + TrilBackwardKernel<<>>( + static_cast(grad_output->DataPtr()), static_cast(grad_input->DataPtr()), rows, cols, + diagonal); + }, + "MACA 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 &maca_stream = dynamic_cast( + infini_train::core::GetDeviceGuardImpl(device.type())->GetStream(device)) + ->maca_stream(); + + core::maca::DispatchMacaFunc( + input->Dtype(), + [=]() { + TriuForwardKernel<<>>( + static_cast(input->DataPtr()), static_cast(output->DataPtr()), rows, cols, diagonal); + }, + "MACA 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 &maca_stream = dynamic_cast( + infini_train::core::GetDeviceGuardImpl(device.type())->GetStream(device)) + ->maca_stream(); + + core::maca::DispatchMacaFunc( + dtype, + [=]() { + grad_input->Fill(0); + TriuBackwardKernel<<>>( + static_cast(grad_output->DataPtr()), static_cast(grad_input->DataPtr()), rows, cols, + diagonal); + }, + "MACA 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)) + ->maca_stream(); + + // Allocate device memory for dims and strides + // TODO(zbl): avoid using mcMalloc? + int64_t *device_buffer; + mcMallocAsync(reinterpret_cast(&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()); + + mcMemcpyAsync(device_buffer, host_buffer.data(), 3 * ndim * sizeof(int64_t), mcMemcpyHostToDevice, stream); + + int threads_per_block = 256; + int num_blocks = (num_elements + threads_per_block - 1) / threads_per_block; + + core::maca::DispatchMacaFunc( + dtype, + [=]() { + output->Fill(0); + TransposeForwardKernel<<>>( + static_cast(input->DataPtr()), static_cast(output->DataPtr()), in_dims_dev, + in_strides_dev, out_strides_dev, ndim, dim0, dim1, num_elements); + }, + "MACA TransposeForward"); + + mcFreeAsync(device_buffer, stream); + + 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 &maca_stream = dynamic_cast( + infini_train::core::GetDeviceGuardImpl(device.type())->GetStream(device)) + ->maca_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::maca::DispatchMacaFunc( + dtype, + [=]() { + MaskLeadsForwardKernel<<>>( + static_cast(input->DataPtr()), static_cast(mask_casted->DataPtr()), + static_cast(output->DataPtr()), common::maca::Cast(value), rows, inner); + }, + "MACA 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::maca::DispatchMacaFunc( + dtype, + [=]() { + MaskForwardKernel<<>>( + static_cast(input->DataPtr()), static_cast(mask_casted->DataPtr()), + static_cast(output->DataPtr()), common::maca::Cast(value), static_cast(batch_size), + static_cast(mask_size)); + }, + "MACA 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 &maca_stream = dynamic_cast( + infini_train::core::GetDeviceGuardImpl(device.type())->GetStream(device)) + ->maca_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::maca::DispatchMacaFunc( + dtype, + [=]() { + grad_input->Fill(0); + MaskLeadsBackwardKernel<<>>( + static_cast(grad_output->DataPtr()), static_cast(mask_casted->DataPtr()), + static_cast(grad_input->DataPtr()), rows, inner); + }, + "MACA 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::maca::DispatchMacaFunc( + dtype, + [=]() { + grad_input->Fill(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)); + }, + "MACA 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 &maca_stream = dynamic_cast( + infini_train::core::GetDeviceGuardImpl(device.type())->GetStream(device)) + ->maca_stream(); + + core::maca::DispatchMacaFunc( + input->Dtype(), + [=]() { + RepeatInterleaveForwardKernel<<>>( + static_cast(input->DataPtr()), static_cast(output->DataPtr()), outer, dim_size, inner, + repeat); + }, + "MACA 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 &maca_stream = dynamic_cast( + infini_train::core::GetDeviceGuardImpl(device.type())->GetStream(device)) + ->maca_stream(); + + core::maca::DispatchMacaFunc( + grad_output->Dtype(), + [=]() { + grad_input->Fill(0); + RepeatInterleaveBackwardKernel<<>>( + static_cast(grad_output->DataPtr()), static_cast(grad_input->DataPtr()), outer, + dim_size, inner, repeat); + }, + "MACA RepeatInterleaveBackward"); + + return grad_input; +} +} // namespace infini_train::kernels::maca + +namespace infini_train::kernels::maca { + +void RegisterTransformKernels() { + REGISTER_KERNEL(infini_train::Device::DeviceType::kPrivateUse1, TrilForward, + infini_train::kernels::maca::TrilForward) + REGISTER_KERNEL(infini_train::Device::DeviceType::kPrivateUse1, TrilBackward, + infini_train::kernels::maca::TrilBackward) + REGISTER_KERNEL(infini_train::Device::DeviceType::kPrivateUse1, TriuForward, + infini_train::kernels::maca::TriuForward) + REGISTER_KERNEL(infini_train::Device::DeviceType::kPrivateUse1, TriuBackward, + infini_train::kernels::maca::TriuBackward) + REGISTER_KERNEL(infini_train::Device::DeviceType::kPrivateUse1, TransposeForward, + infini_train::kernels::maca::TransposeForward) + REGISTER_KERNEL(infini_train::Device::DeviceType::kPrivateUse1, TransposeBackward, + infini_train::kernels::maca::TransposeBackward) + REGISTER_KERNEL(infini_train::Device::DeviceType::kPrivateUse1, MaskForward, + infini_train::kernels::maca::MaskForward) + REGISTER_KERNEL(infini_train::Device::DeviceType::kPrivateUse1, MaskBackward, + infini_train::kernels::maca::MaskBackward) + REGISTER_KERNEL(infini_train::Device::DeviceType::kPrivateUse1, RepeatInterleaveForward, + infini_train::kernels::maca::RepeatInterleaveForward) + REGISTER_KERNEL(infini_train::Device::DeviceType::kPrivateUse1, RepeatInterleaveBackward, + infini_train::kernels::maca::RepeatInterleaveBackward) +} + +} // namespace infini_train::kernels::maca diff --git a/backends/maca/src/kernels/vocab_parallel_cross_entropy.maca b/backends/maca/src/kernels/vocab_parallel_cross_entropy.maca new file mode 100644 index 0000000..6581ec0 --- /dev/null +++ b/backends/maca/src/kernels/vocab_parallel_cross_entropy.maca @@ -0,0 +1,132 @@ +#include +#include +#include +#include + +#include + +#include "infini_train/include/core/runtime/device_guard.h" +#include "infini_train/include/dispatcher.h" +#include "infini_train/include/tensor.h" + +#include "common/common_maca.h" +#include "common/kernel_helper.cuh" +#include "runtime/maca_dispatch.h" +#include "runtime/maca_runtime_common.h" + +namespace infini_train::kernels::maca { + +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::maca::Cast(dloss_is_scalar ? dloss_buf[0] : dloss_buf[r]); + const float vm_row = 1.0f - common::maca::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::maca::Cast(softmax_local[idx]); + const float vm = common::maca::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::maca::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 &maca_stream = dynamic_cast( + infini_train::core::GetDeviceGuardImpl(device.type())->GetStream(device)) + ->maca_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::maca::DispatchMacaFunc, + 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); + }, + "MACA VocabParallelCrossEntropyBackward"); + + return grad_input; +} +} // namespace infini_train::kernels::maca + +namespace infini_train::kernels::maca { + +void RegisterVocabParallelCrossEntropyKernels() { + REGISTER_KERNEL(infini_train::Device::DeviceType::kPrivateUse1, VocabParallelCrossEntropyBackward, + infini_train::kernels::maca::VocabParallelCrossEntropyBackward) +} + +} // namespace infini_train::kernels::maca diff --git a/backends/maca/src/runtime/maca_dispatch.h b/backends/maca/src/runtime/maca_dispatch.h new file mode 100644 index 0000000..791bf4e --- /dev/null +++ b/backends/maca/src/runtime/maca_dispatch.h @@ -0,0 +1,53 @@ +#pragma once + +#include +#include +#include + +#include +#include + +#include "infini_train/include/core/backend_type_map.h" +#include "infini_train/include/dtype_dispatch.h" + +// ----------------------------------------------------------------------------- +// MACA low-precision BackendTypeMap specializations: +// FP16 -> __half, BF16 -> __maca_bfloat16 +// ----------------------------------------------------------------------------- +namespace infini_train::core { +template <> struct BackendTypeMap { + using type = __half; +}; + +template <> struct BackendTypeMap { + using type = __maca_bfloat16; +}; +} // namespace infini_train::core + +// Register all standard (non-low-precision) dtypes for the MACA backend. +// FP16/BF16 are registered explicitly above with their MACA-native scalar +// types. +INFINI_REGISTER_STANDARD_BACKEND_TYPES(infini_train::Device::DeviceType::kPrivateUse1) + +namespace infini_train::core::maca { + +template struct MacaTypeMap : BackendTypeMap {}; + +// ----------------------------------------------------------------------------- +// MACA dispatch helpers +// ----------------------------------------------------------------------------- + +template +auto DispatchMacaFunc(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 DispatchMacaFunc(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::maca diff --git a/backends/maca/src/runtime/maca_guard_impl.cc b/backends/maca/src/runtime/maca_guard_impl.cc new file mode 100644 index 0000000..e48bab8 --- /dev/null +++ b/backends/maca/src/runtime/maca_guard_impl.cc @@ -0,0 +1,369 @@ +#include "runtime/maca_guard_impl.h" + +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include + +#include "infini_train/include/core/runtime/runtime_common.h" +#include "infini_train/include/device.h" + +#include "common/common_maca.h" +#include "runtime/maca_runtime_common.h" + +namespace infini_train::core::maca { +namespace { +// Read /proc/self/cmdline and return --tensor_parallel value, or 1 if absent +// or unparseable. This does not depend on gflags and runs before mcInit. +int ReadTensorParallelFromCmdline() { + std::ifstream in("/proc/self/cmdline", std::ios::binary); + if (!in) { + return 1; + } + + std::vector args; + std::string current; + char c; + while (in.get(c)) { + if (c == '\0') { + if (!current.empty()) { + args.push_back(std::move(current)); + current.clear(); + } + } else { + current.push_back(c); + } + } + if (!current.empty()) { + args.push_back(std::move(current)); + } + + constexpr char kTensorParallelFlag[] = "--tensor_parallel"; + constexpr char kTensorParallelPrefix[] = "--tensor_parallel="; + for (size_t i = 0; i < args.size(); ++i) { + std::string value; + if (args[i].rfind(kTensorParallelPrefix, 0) == 0) { + value = args[i].substr(sizeof(kTensorParallelPrefix) - 1); + } else if (args[i] == kTensorParallelFlag && i + 1 < args.size()) { + value = args[i + 1]; + } else { + continue; + } + + try { + return std::stoi(value); + } catch (...) { return 1; } + } + return 1; +} + +static std::vector> g_maca_streams; +static std::vector> g_maca_blas_handles; +static std::vector> g_device_stream_flags; +static std::vector> g_device_handle_flags; + +// Serialize host-side MemcpyAsync across threads. On MACA, concurrent +// mcMemcpyAsync from multiple threads during init-time bursts +// (Module::To uploads, Adam state fills, ...) races with the runtime's +// auto P2P peer-mapping and produces "readonly page" faults or +// mcErrorInvalidValue. The lock is held only for the brief window of the +// API call itself; actual GPU work remains async on the caller's stream. +static std::mutex g_memcpy_mutex; + +inline void CheckMacaDevice(Device device) { + CHECK(device.type() == Device::DeviceType::kPrivateUse1) << std::format( + "MacaGuardImpl expects MACA device, but got type={} index={}", static_cast(device.type()), device.index()); + const int idx = device.index(); + CHECK(idx >= 0 && static_cast(idx) < g_maca_streams.size()) + << std::format("MACA device index {} out of cache range [0, {}).", idx, g_maca_streams.size()); +} + +inline mcEvent_t GetMacaEvent(Event *event) { + auto *maca_event = dynamic_cast(event); + CHECK_NOTNULL(maca_event); + return maca_event->maca_event(); +} + +inline mcStream_t GetMacaStream(Stream *stream) { + auto *maca_stream = dynamic_cast(stream); + CHECK_NOTNULL(maca_stream); + return maca_stream->maca_stream(); +} +} // namespace + +void MacaGuardImpl::InitSingleStream(Device device) { + CheckMacaDevice(device); + + int current_device = -1; + MACA_CHECK(mcGetDevice(¤t_device)); + MACA_CHECK(mcSetDevice(device.index())); + + g_maca_streams[device.index()] = std::make_unique(); + + MACA_CHECK(mcSetDevice(current_device)); +} + +void MacaGuardImpl::InitSingleHandle(Device device) { + CheckMacaDevice(device); + + int current_device = -1; + MACA_CHECK(mcGetDevice(¤t_device)); + MACA_CHECK(mcSetDevice(device.index())); + + std::call_once(*g_device_stream_flags.at(device.index()), InitSingleStream, device); + + g_maca_blas_handles[device.index()] = std::make_unique(g_maca_streams[device.index()].get()); + + MACA_CHECK(mcSetDevice(current_device)); +} + +void MacaGuardImpl::Initialize() { + std::call_once(initialize_flag_, [] { + // FIXME(cx): Stop deriving runtime policy from argv and mutating process-wide + // environment here. Pass MACA runtime/communication policy through explicit + // provider or launcher configuration instead. + // Apply provider runtime policy immediately before mcInit. Users may + // override it in the process environment before first device use. + setenv("MACA_LAUNCH_BLOCKING", "1", 0); + if (ReadTensorParallelFromCmdline() > 1) { + setenv("MCCL_P2P_DISABLE", "1", 0); + } + MACA_CHECK(mcInit(0)); + + int device_count = 0; + MACA_CHECK(mcGetDeviceCount(&device_count)); + CHECK_GT(device_count, 0) << "No MACA devices are available."; + CHECK_LE(static_cast(device_count), static_cast(std::numeric_limits::max()) + 1) + << "MACA device count exceeds InfiniTrain's Device index range."; + + g_maca_streams.resize(device_count); + g_maca_blas_handles.resize(device_count); + g_device_stream_flags.reserve(device_count); + g_device_handle_flags.reserve(device_count); + for (int i = 0; i < device_count; ++i) { + g_device_stream_flags.push_back(std::make_unique()); + g_device_handle_flags.push_back(std::make_unique()); + } + }); +} + +// device +Device MacaGuardImpl::GetDevice() const { + int current_device = -1; + MACA_CHECK(mcGetDevice(¤t_device)); + return Device(Device::DeviceType::kPrivateUse1, current_device); +} + +void MacaGuardImpl::SetDevice(Device device) const { + CheckMacaDevice(device); + MACA_CHECK(mcSetDevice(device.index())); +} + +int MacaGuardImpl::DeviceCount() const { return static_cast(g_maca_streams.size()); } + +Device::DeviceType MacaGuardImpl::Type() const { return Device::DeviceType::kPrivateUse1; } + +// stream +Stream *MacaGuardImpl::GetStream(Device device) const { + CheckMacaDevice(device); + std::call_once(*g_device_stream_flags.at(device.index()), InitSingleStream, device); + return g_maca_streams.at(device.index()).get(); +} + +Stream *MacaGuardImpl::CreateStream(Device device) const { + CheckMacaDevice(device); + int current_device = -1; + MACA_CHECK(mcGetDevice(¤t_device)); + MACA_CHECK(mcSetDevice(device.index())); + + Stream *stream = new MacaStream(); + + MACA_CHECK(mcSetDevice(current_device)); + return stream; +} + +Stream *MacaGuardImpl::CreateStreamWithPriority(Device device, int priority) const { + CheckMacaDevice(device); + int current_device = -1; + MACA_CHECK(mcGetDevice(¤t_device)); + MACA_CHECK(mcSetDevice(device.index())); + + Stream *stream = new MacaStream(priority); + + MACA_CHECK(mcSetDevice(current_device)); + return stream; +} + +void MacaGuardImpl::DestroyStream(Stream *stream) const { + if (stream == nullptr) { + return; + } + auto *maca_stream = dynamic_cast(stream); + CHECK_NOTNULL(maca_stream); + MACA_CHECK(mcStreamDestroy(maca_stream->maca_stream())); + delete maca_stream; +} + +void MacaGuardImpl::GetStreamPriorityRange(int *low, int *high) const { + MACA_CHECK(mcDeviceGetStreamPriorityRange(low, high)); +} + +// event +void MacaGuardImpl::EventCreate(Event **event) const { *event = new MacaEvent(); } + +void MacaGuardImpl::EventCreateWithFlags(Event **event, EventFlag flags) const { *event = new MacaEvent(flags); } + +void MacaGuardImpl::EventDestroy(Event *event) const { + if (event == nullptr) { + return; + } + delete event; +} + +void MacaGuardImpl::EventRecord(Event *event, Stream *stream) const { + auto maca_event = GetMacaEvent(event); + auto maca_stream = GetMacaStream(stream); + MACA_CHECK(mcEventRecord(maca_event, maca_stream)); +} + +void MacaGuardImpl::StreamWaitEvent(Stream *stream, Event *event, uint32_t flags) const { + auto maca_event = GetMacaEvent(event); + auto maca_stream = GetMacaStream(stream); + MACA_CHECK(mcStreamWaitEvent(maca_stream, maca_event, flags)); +} + +RuntimeStatus MacaGuardImpl::EventSynchronize(Event *event) const { + auto maca_event = GetMacaEvent(event); + mcError_t status = mcEventSynchronize(maca_event); + if (status == mcSuccess) { + return RuntimeStatus::kSuccess; + } + if (status == mcErrorNotReady) { + return RuntimeStatus::kNotReady; + } + LOG(ERROR) << "MacaGuardImpl::EventSynchronize failed: " << mcGetErrorString(status); + return RuntimeStatus::kError; +} + +RuntimeStatus MacaGuardImpl::EventQuery(Event *event) const { + auto maca_event = GetMacaEvent(event); + mcError_t status = mcEventQuery(maca_event); + if (status == mcSuccess) { + return RuntimeStatus::kSuccess; + } + if (status == mcErrorNotReady) { + return RuntimeStatus::kNotReady; + } + LOG(ERROR) << "MacaGuardImpl::EventQuery failed: " << mcGetErrorString(status); + return RuntimeStatus::kError; +} + +float MacaGuardImpl::EventElapsedTime(Event *start_event, Event *stop_event) const { + auto start_maca_event = GetMacaEvent(start_event); + auto stop_maca_event = GetMacaEvent(stop_event); + float elapsed_ms = 0.0f; + MACA_CHECK(mcEventElapsedTime(&elapsed_ms, start_maca_event, stop_maca_event)); + return elapsed_ms; +} + +// sync +void MacaGuardImpl::SynchronizeDevice(Device device) const { + auto original_device = GetDevice(); + SetDevice(device); + + MACA_CHECK(mcDeviceSynchronize()); + + SetDevice(original_device); +} + +void MacaGuardImpl::SynchronizeStream(Stream *stream) const { + auto maca_stream = GetMacaStream(stream); + MACA_CHECK(mcStreamSynchronize(maca_stream)); +} + +// blas +BlasHandle *MacaGuardImpl::GetBlasHandle(Device device) const { + CheckMacaDevice(device); + std::call_once(*g_device_handle_flags.at(device.index()), InitSingleHandle, device); + return g_maca_blas_handles.at(device.index()).get(); +} + +// memory +void MacaGuardImpl::Malloc(void **dev_ptr, size_t size) { MACA_CHECK(mcMalloc(dev_ptr, size)); } + +void MacaGuardImpl::MallocAsync(void **dev_ptr, size_t size, Stream *stream) { + // NOTE(dcj): mcMallocAsync with a per-stream mempool gives a big speedup + // (~2x on gpt2 DDP steady-state) vs synchronous mcMalloc, but under + // multi-thread DDP init bursts (e.g. llama3 1B with nthread=8 uploading + // hundreds of param tensors) it races with MACA's auto P2P peer-mapping + // and produces mcErrorInvalidValue on subsequent mcMemcpyAsync, or + // "readonly page" faults -- no amount of mutex/stream-sync serialization + // around the alloc call suppresses this. Keep the synchronous path for + // correctness. + // auto maca_stream = GetMacaStream(stream); + // MACA_CHECK(mcMallocAsync(dev_ptr, size, maca_stream)); + (void)stream; + Malloc(dev_ptr, size); +} + +void MacaGuardImpl::Free(void *dev_ptr) { MACA_CHECK(mcFree(dev_ptr)); } + +void MacaGuardImpl::FreeAsync(void *dev_ptr, Stream *stream) { + // auto maca_stream = GetMacaStream(stream); + // MACA_CHECK(mcFreeAsync(dev_ptr, maca_stream)); + (void)stream; + Free(dev_ptr); +} + +void MacaGuardImpl::Memcpy(void *dst, const void *src, size_t count, MemcpyKind kind) { + if (kind == MemcpyKind::kH2D) { + MACA_CHECK(mcMemcpy(dst, src, count, mcMemcpyHostToDevice)); + } else if (kind == MemcpyKind::kD2H) { + MACA_CHECK(mcMemcpy(dst, src, count, mcMemcpyDeviceToHost)); + } else if (kind == MemcpyKind::kD2D) { + MACA_CHECK(mcMemcpy(dst, src, count, mcMemcpyDeviceToDevice)); + } else { + LOG(FATAL) << std::format("MacaGuardImpl::Memcpy got invalid MemcpyKind={}", MemcpyKindToString(kind)); + } +} + +void MacaGuardImpl::MemcpyAsync(void *dst, const void *src, size_t count, MemcpyKind kind, Stream *stream) { + std::lock_guard lock(g_memcpy_mutex); + auto maca_stream = GetMacaStream(stream); + + switch (kind) { + case MemcpyKind::kH2D: + MACA_CHECK(mcMemcpyAsync(dst, src, count, mcMemcpyHostToDevice, maca_stream)); + break; + case MemcpyKind::kD2H: + MACA_CHECK(mcMemcpyAsync(dst, src, count, mcMemcpyDeviceToHost, maca_stream)); + break; + case MemcpyKind::kD2D: + MACA_CHECK(mcMemcpyAsync(dst, src, count, mcMemcpyDeviceToDevice, maca_stream)); + break; + default: + LOG(FATAL) << std::format("MacaGuardImpl::MemcpyAsync got invalid MemcpyKind={}", MemcpyKindToString(kind)); + } +} + +void MacaGuardImpl::ResetMemPoolHighWatermarks(Device device) const { + // MetaX SDK support for the mempool high-watermark attributes is not + // confirmed. Keep this a no-op, matching feat/muxi_device_registry. + (void)device; +} + +std::pair MacaGuardImpl::GetMemPoolPeakMB(Device device) const { + (void)device; + return std::make_pair(0, 0); +} + +void RegisterMacaRuntime() { INFINI_TRAIN_REGISTER_DEVICE_GUARD_IMPL(Device::DeviceType::kPrivateUse1, MacaGuardImpl) } + +} // namespace infini_train::core::maca diff --git a/backends/maca/src/runtime/maca_guard_impl.h b/backends/maca/src/runtime/maca_guard_impl.h new file mode 100644 index 0000000..4420e21 --- /dev/null +++ b/backends/maca/src/runtime/maca_guard_impl.h @@ -0,0 +1,93 @@ +#pragma once + +#include +#include +#include +#include + +#include "infini_train/include/core/runtime/device_guard.h" + +namespace infini_train::core { +class Stream; +class BlasHandle; +} // namespace infini_train::core + +namespace infini_train::core::maca { + +void RegisterMacaRuntime(); + +class MacaGuardImpl final : public DeviceGuardImpl { +public: + static void InitSingleStream(Device device); + + static void InitSingleHandle(Device device); + + MacaGuardImpl() = default; + void Initialize() override; + + // device + Device GetDevice() const override; + + void SetDevice(Device device) const override; + + int DeviceCount() const override; + + Device::DeviceType Type() const override; + + // stream + Stream *GetStream(Device device) const override; + + Stream *CreateStream(Device device) const override; + + Stream *CreateStreamWithPriority(Device device, int priority) const override; + + void DestroyStream(Stream *stream) const override; + + void GetStreamPriorityRange(int *low, int *high) const override; + + // event + void EventCreate(Event **event) const override; + + void EventCreateWithFlags(Event **event, EventFlag flags) const override; + + void EventDestroy(Event *event) const override; + + void EventRecord(Event *event, Stream *stream) const override; + + void StreamWaitEvent(Stream *stream, Event *event, uint32_t flags) const override; + + RuntimeStatus EventSynchronize(Event *event) const override; + + RuntimeStatus EventQuery(Event *event) const override; + + float EventElapsedTime(Event *start_event, Event *stop_event) const override; + + // sync + void SynchronizeDevice(Device device) const override; + void SynchronizeStream(Stream *stream) const override; + + // blas + BlasHandle *GetBlasHandle(Device device) const override; + + // memory + void Malloc(void **dev_ptr, size_t size) override; + + void MallocAsync(void **dev_ptr, size_t size, Stream *stream) override; + + void Free(void *dev_ptr) override; + + void FreeAsync(void *dev_ptr, Stream *stream) override; + + void Memcpy(void *dst, const void *src, size_t count, MemcpyKind kind) override; + + void MemcpyAsync(void *dst, const void *src, size_t count, MemcpyKind kind, Stream *stream) override; + + void ResetMemPoolHighWatermarks(Device device) const override; + + std::pair GetMemPoolPeakMB(Device device) const override; + +private: + std::once_flag initialize_flag_; +}; + +} // namespace infini_train::core::maca diff --git a/backends/maca/src/runtime/maca_runtime_common.cc b/backends/maca/src/runtime/maca_runtime_common.cc new file mode 100644 index 0000000..6a178a5 --- /dev/null +++ b/backends/maca/src/runtime/maca_runtime_common.cc @@ -0,0 +1,60 @@ +#include "runtime/maca_runtime_common.h" + +#include "common/common_maca.h" + +namespace infini_train::core::maca { +namespace { +uint32_t ToMacaEventFlags(EventFlag flags) { + switch (flags) { + case EventFlag::kDefault: + return mcEventDefault; + case EventFlag::kBlockingSync: + return mcEventBlockingSync; + case EventFlag::kDisableTiming: + return mcEventDisableTiming; + case EventFlag::kInterprocess: + // MACA (like CUDA) requires DisableTiming with Interprocess events. + // NOTE(dcj): if the MACA SDK in use does not expose mcEventInterprocess, + // this branch will need to be guarded and downgraded to a LOG(FATAL). + return mcEventInterprocess | mcEventDisableTiming; + default: + LOG(FATAL) << "Unsupported EventFlag value: " << static_cast(flags); + } + return mcEventDefault; +} +} // namespace + +MacaEvent::MacaEvent(EventFlag flags) { MACA_CHECK(mcEventCreateWithFlags(&event_, ToMacaEventFlags(flags))); } + +MacaEvent::~MacaEvent() { + if (event_ != nullptr) { + MACA_CHECK(mcEventDestroy(event_)); + } +} + +mcEvent_t MacaEvent::maca_event() const { return event_; } + +MacaStream::MacaStream() { MACA_CHECK(mcStreamCreate(&stream_)); } + +MacaStream::MacaStream(int priority) { + MACA_CHECK(mcStreamCreateWithPriority(&stream_, mcStreamNonBlocking, priority)); +} + +MacaStream::~MacaStream() { + // Do nothing. +} + +mcStream_t MacaStream::maca_stream() const { return stream_; } + +MacaBlasHandle::MacaBlasHandle(Stream *stream) { + MCBLAS_CHECK(mcblasCreate(&mcblas_handle_)); + MCBLAS_CHECK(mcblasSetStream(mcblas_handle_, dynamic_cast(stream)->maca_stream())); +} + +MacaBlasHandle::~MacaBlasHandle() { + // Do nothing. +} + +mcblasHandle_t MacaBlasHandle::mcblas_handle() const { return mcblas_handle_; } + +} // namespace infini_train::core::maca diff --git a/backends/maca/src/runtime/maca_runtime_common.h b/backends/maca/src/runtime/maca_runtime_common.h new file mode 100644 index 0000000..fb4b0a8 --- /dev/null +++ b/backends/maca/src/runtime/maca_runtime_common.h @@ -0,0 +1,59 @@ +#pragma once + +#include + +#include +#include +#include + +#include "infini_train/include/core/runtime/runtime_common.h" + +namespace infini_train::core { +class Stream; +} + +namespace infini_train::core::maca { + +class MacaEvent final : public Event { +public: + explicit MacaEvent(EventFlag flags = EventFlag::kDefault); + ~MacaEvent() override; + + mcEvent_t maca_event() const; + +private: + mcEvent_t event_ = nullptr; +}; + +class MacaStream : public Stream { +public: + MacaStream(); + explicit MacaStream(int priority); + + // NOTE(dcj): + // The wrapper destructor avoids MACA calls because global destruction may + // outlive the runtime. Explicitly owned streams must be released through + // DeviceGuardImpl::DestroyStream while the runtime is alive. + ~MacaStream() override; + + mcStream_t maca_stream() const; + +private: + mcStream_t stream_ = nullptr; +}; + +class MacaBlasHandle : public BlasHandle { +public: + explicit MacaBlasHandle(Stream *stream); + + // NOTE(dcj): + // Mirror CudaBlasHandle: leaked intentionally; see MacaStream note. + ~MacaBlasHandle() override; + + mcblasHandle_t mcblas_handle() const; + +private: + mcblasHandle_t mcblas_handle_; +}; + +} // namespace infini_train::core::maca diff --git a/scripts/format.py b/scripts/format.py new file mode 100644 index 0000000..e38b2b3 --- /dev/null +++ b/scripts/format.py @@ -0,0 +1,207 @@ +import argparse +import subprocess +import os +from pathlib import Path +from colorama import Fore, Style + +# Supported file types +SUPPORTED_FILES = { + ".h": "c", + ".hh": "c", + ".hpp": "c", + ".c": "c", + ".cc": "c", + ".cpp": "c", + ".cxx": "c", + ".cu": "c", + ".cuh": "c", + ".maca": "c", + ".mlu": "c", + ".cl": "c", + ".py": "py", +} + + +def format_file(file: Path, check: bool, formatter) -> bool: + formatter = formatter.get(SUPPORTED_FILES.get(file.suffix, None), None) + if not formatter: + return True # File type not supported, skip + + try: + cmd = [] + if formatter.startswith("clang-format"): + cmd = [formatter, "-style=file", "-i", file] + if check: + cmd.insert(2, "-dry-run") + process = subprocess.run( + cmd, + capture_output=True, + text=True, + check=True, + ) + if process.stderr: + print(f"{Fore.YELLOW}{file} is not formatted.{Style.RESET_ALL}") + print( + f"Use {Fore.CYAN}{formatter} -style=file -i {file}{Style.RESET_ALL} to format it." + ) + return False + else: + subprocess.run( + cmd, + capture_output=True, + text=True, + check=True, + ) + print(f"{Fore.CYAN}Formatted: {file}{Style.RESET_ALL}") + elif formatter == "black": + cmd = [formatter, file] + if check: + cmd.insert(1, "--check") + process = subprocess.run( + cmd, + capture_output=True, + text=True, + check=True, + ) + if process.stderr: + print(f"{Fore.YELLOW}{file} is not formatted.{Style.RESET_ALL}") + print( + f"Use {Fore.CYAN}{formatter} {file}{Style.RESET_ALL} to format it." + ) + return False + else: + subprocess.run( + cmd, + capture_output=True, + text=True, + check=True, + ) + print(f"{Fore.CYAN}Formatted: {file}{Style.RESET_ALL}") + except FileNotFoundError: + print( + f"{Fore.RED}Formatter {formatter} not found.{Style.RESET_ALL}" + ) + return False + except subprocess.CalledProcessError as e: + print(f"{Fore.RED}Formatter {formatter} failed: {e}{Style.RESET_ALL}") + + return True + + +def git_added_files(): + """Get all files with staged changes""" + try: + # Use git diff --cached --name-only to get all files added to the staging area + result = subprocess.run( + ["git", "diff", "--cached", "--diff-filter=AMR", "--name-only"], + capture_output=True, + text=True, + check=True, + ) + for file in result.stdout.splitlines(): + yield Path(file.strip()) + except subprocess.CalledProcessError as e: + print(f"{Fore.RED}Git diff failed: {e}{Style.RESET_ALL}") + + +def git_modified_since_ref(ref): + """Get a list of modified files from the specified Git reference to the current state""" + try: + result = subprocess.run( + ["git", "diff", f"{ref}..", "--diff-filter=AMR", "--name-only"], + capture_output=True, + text=True, + check=True, + ) + for file in result.stdout.splitlines(): + yield Path(file.strip()) + except subprocess.CalledProcessError as e: + print(f"{Fore.RED}Git diff failed: {e}{Style.RESET_ALL}") + + +def list_files(paths): + """Recursively get all files under the specified path""" + files = [] + for path in paths: + if path.is_file(): + yield path + elif path.is_dir(): + for dirpath, _, filenames in os.walk(path): + for name in filenames: + yield Path(dirpath) / name + else: + print( + f"{Fore.RED}Error: {path} is not a file or directory.{Style.RESET_ALL}" + ) + + +def filter_in_path(file: Path, path) -> bool: + """Determine whether the file is in the specified path""" + for p in path: + if file.is_relative_to(p): + return True + return False + + +def main(): + parser = argparse.ArgumentParser() + parser.add_argument( + "--ref", type=str, help="Git reference (commit hash) to compare against." + ) + parser.add_argument( + "--path", nargs="*", type=Path, help="Files to format or check." + ) + parser.add_argument( + "--check", action="store_true", help="Check files without modifying them." + ) + parser.add_argument( + "--c", default="clang-format-16", help="C formatter (default: clang-format-16)" + ) + parser.add_argument( + "--py", default="black", help="Python formatter (default: black)" + ) + args = parser.parse_args() + + if args.ref is None and args.path is None: + # Last commit. + print(f"{Fore.GREEN}Formating git added files.{Style.RESET_ALL}") + files = git_added_files() + + else: + if args.ref is None: + print(f"{Fore.GREEN}Formating files in {args.path}.{Style.RESET_ALL}") + files = list_files(args.path) + elif args.path is None: + print( + f"{Fore.GREEN}Formating git modified files from {args.ref}.{Style.RESET_ALL}" + ) + files = git_modified_since_ref(args.ref) + else: + print( + f"{Fore.GREEN}Formating git modified files from {args.ref} in {args.path}.{Style.RESET_ALL}" + ) + files = ( + file + for file in git_modified_since_ref(args.ref) + if filter_in_path(file, args.path) + ) + + formatted = True + for file in files: + if not format_file( + file, + args.check, + { + "c": args.c, + "py": args.py, + }, + ): + formatted = False + + if not formatted: + exit(1) + + +if __name__ == "__main__": + main() + diff --git a/third_party/InfiniTrain b/third_party/InfiniTrain new file mode 160000 index 0000000..4b1042c --- /dev/null +++ b/third_party/InfiniTrain @@ -0,0 +1 @@ +Subproject commit 4b1042cfb9ec629c44989239ced3be12cab4f81e