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/.clang-tidy b/.clang-tidy new file mode 100644 index 0000000..65c7517 --- /dev/null +++ b/.clang-tidy @@ -0,0 +1,17 @@ +--- +Checks: > + -*, + infinitrain-no-member-definitions-in-headers, + readability-identifier-naming +WarningsAsErrors: '*' +HeaderFilterRegex: '(^|.*/)(backends|tools)/' +HeaderFileExtensions: [h, hh, hpp, hxx, cuh] +FormatStyle: none +CheckOptions: + # InfiniTrain custom checks. + infinitrain-no-member-definitions-in-headers.IgnoreEmptyBodies: true + + # Enum constants: kFooBar + readability-identifier-naming.EnumConstantCase: CamelCase + readability-identifier-naming.EnumConstantPrefix: k +... diff --git a/.git-blame-ignore-revs b/.git-blame-ignore-revs new file mode 100644 index 0000000..38781ba --- /dev/null +++ b/.git-blame-ignore-revs @@ -0,0 +1,2 @@ +# Formatting-only commits; preserve attribution to the original changes. +baddcda9429658cb0d4a55e850302940a58b868f diff --git a/.github/workflows/format-check.yaml b/.github/workflows/format-check.yaml new file mode 100644 index 0000000..15e4047 --- /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 formatters + run: | + python3 -m pip install --upgrade pip + python3 -m pip install black clang-format==16.0.6 + + - name: Run format check + run: | + python3 scripts/format.py --c clang-format --path backends --check + diff --git a/.gitignore b/.gitignore new file mode 100644 index 0000000..eeba729 --- /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..9220950 100644 --- a/README.md +++ b/README.md @@ -1 +1,119 @@ -# InfiniTrain-Backends \ No newline at end of file +# InfiniTrain Backends + +**English** | [简体中文](README.zh-CN.md) + +[![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 the +InfiniTrain submodule pinned to `v0.7.0`. 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](backends/maca/README.md) | MACA | MCCL (optional) | GPT-2, LLaMA 3, Mixtral | + +The model sources normally come from the pinned InfiniTrain submodule. Each +backend supplies the runtime, kernels, collective implementation, and build +integration needed to run its supported examples on the corresponding +accelerator. See the linked backend guides for setup and usage. + +## Common Requirements + +- **Linux** +- **CMake**: 3.28 or newer +- **Git** with submodule support +- A **C++20** compiler and vendor SDK compatible with the selected backend + +Additional dependencies and SDK setup are documented in each backend's guide. + +## Build Options + +These options apply across backends. Provider-specific options are documented +in the corresponding backend guide. + +| 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 unit tests and variants for the selected provider | + +Use a separate build directory for each provider so compiler and SDK cache +entries do not leak between them. Top-level builds enable backend examples by +default. + +For development against another InfiniTrain checkout, add +`-DINFINITRAIN_SOURCE_DIR=/path/to/InfiniTrain` to the configure command. The +selected checkout must implement the PrivateUse1 extension API expected by +the backend. + +PrivateUse1 providers require `USE_CUDA=OFF`; configuration fails rather than +silently overriding an explicit `USE_CUDA=ON`. Unit tests for the selected +provider use the PrivateUse1 device type and `PRIVATEUSE1/*` GTest prefix, +independent of the provider's name. + +## Architecture + +```text +backends// + README.md provider setup and usage guide + cmake/ compiler and vendor SDK setup before project() + include/ public provider API + src/backend.cc provider metadata 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`. + +When adding this repository as a third-party library with `add_subdirectory()`, +select the compiler required by the provider in the parent project before +enabling the C++ language. After setting up the vendor SDK environment, use +the following order in the parent's `CMakeLists.txt`, replacing the compiler +path and provider name: + +```cmake +cmake_minimum_required(VERSION 3.28) + +set(CMAKE_CXX_COMPILER "/path/to/vendor/compiler" CACHE FILEPATH "") +set(INFINITRAIN_BACKEND "" CACHE STRING "") + +project(MyTrainingApp LANGUAGES CXX) +add_subdirectory(third_party/InfiniTrain-Backends) +``` + +CMake cannot replace a compiler after a language has been enabled. + +## Development checks + +See [formatting and static checks](scripts/README.md) for clang-format, +clang-tidy and IWYU setup, scope and manual usage (in Chinese). + +## License + +InfiniTrain Backends is released under the [MIT License](LICENSE). diff --git a/README.zh-CN.md b/README.zh-CN.md new file mode 100644 index 0000000..27b855e --- /dev/null +++ b/README.zh-CN.md @@ -0,0 +1,109 @@ +# InfiniTrain Backends + +[English](README.md) | **简体中文** + +[![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 为 [InfiniTrain](https://github.com/InfiniTensor/InfiniTrain) +提供独立于框架核心维护的加速器后端。各后端实现 InfiniTrain 的 `PrivateUse1` +接口，将厂商 SDK 集成、运行时支持、集合通信和算子实现放在框架核心之外。 + +每个后端位于独立的 `backends/` 目录中，通常使用固定为 `v0.7.0` 的 +InfiniTrain 子模块。每个构建目录只选择一个后端，每个进程最多注册一个 +`DeviceType::kPrivateUse1` 后端。 + +## 支持的后端 + +| 后端 | 运行时 | 集合通信 | 模型示例 | +| ---- | ------ | -------- | -------- | +| [MACA](backends/maca/README.zh-CN.md) | MACA | MCCL（可选） | GPT-2、LLaMA 3、Mixtral | + +模型源码通常来自锁定版本的 InfiniTrain 子模块。各后端提供运行时、算子、集合通信 +实现和构建集成，以便在对应加速器上运行其支持的模型示例。环境配置和使用方式见上表 +链接的各后端指南。 + +## 通用依赖 + +- **Linux** +- **CMake**：3.28 或更新版本 +- 支持子模块的 **Git** +- 与所选后端兼容的 **C++20** 编译器和厂商 SDK + +额外依赖及 SDK 配置方式见对应后端指南。 + +## 构建选项 + +以下选项适用于各后端。后端专用选项见对应后端指南。 + +| 选项 | 默认值 | 说明 | +| ---- | ------ | ---- | +| `INFINITRAIN_BACKEND` | 必填 | 选择 `backends/` 下的后端 | +| `INFINITRAIN_SOURCE_DIR` | `third_party/InfiniTrain` | InfiniTrain 源码目录，通常为锁定版本的子模块 | +| `INFINITRAIN_BACKENDS_BUILD_EXAMPLES` | 独立构建本仓库时为 `ON` | 构建启用所选后端的 InfiniTrain 示例 | +| `BUILD_TEST` | `OFF` | 构建 InfiniTrain 单元测试及所选后端的测试版本 | + +请为不同后端使用独立的构建目录，避免混用编译器和 SDK 缓存。独立构建本仓库时， +默认启用后端示例。 + +如需使用其他 InfiniTrain 工作副本进行开发，可在 CMake 配置命令中添加 +`-DINFINITRAIN_SOURCE_DIR=/path/to/InfiniTrain`。所选源码版本必须实现该后端所需的 +PrivateUse1 扩展接口。 + +PrivateUse1 后端要求 `USE_CUDA=OFF`。如果显式设置 `USE_CUDA=ON`，配置会报错， +而不会静默覆盖该选项。所选后端的单元测试使用 PrivateUse1 设备类型和 +`PRIVATEUSE1/*` GTest 前缀，与后端名称无关。 + +## 架构 + +```text +backends// + README.md 后端环境配置与使用指南 + cmake/ 在 project() 前配置编译器和厂商 SDK + include/ 后端公共接口 + src/backend.cc 后端元数据注册入口 + src/common/ 后端内部共享辅助代码 + src/runtime/ 设备、流、事件和内存分配集成 + src/kernels/ 后端算子实现及注册 + src/ccl/ 可选的集合通信集成 + examples/ 上游 InfiniTrain 示例的 CMake 适配 + scripts/ 后端模型运行配置及封装脚本 + +third_party/InfiniTrain/ 框架源码，通常为锁定版本的 Git 子模块 +``` + +根目录构建脚本在首次调用 `project()` 之前，根据所选后端加载 +`backends//cmake/pre_project.cmake`。各后端可在此选择编译器并准备 +SDK 相关的依赖探测配置，无需在根目录 `CMakeLists.txt` 中增加厂商分支。 + +通过 `add_subdirectory()` 将本仓库作为第三方库添加到父项目时，应在启用 C++ +语言前选择后端所需的编译器。配置好厂商 SDK 环境后，可按以下顺序编写父项目的 +`CMakeLists.txt`，并替换编译器路径和后端名称： + +```cmake +cmake_minimum_required(VERSION 3.28) + +set(CMAKE_CXX_COMPILER "/path/to/vendor/compiler" CACHE FILEPATH "") +set(INFINITRAIN_BACKEND "" CACHE STRING "") + +project(MyTrainingApp LANGUAGES CXX) +add_subdirectory(third_party/InfiniTrain-Backends) +``` + +CMake 无法在启用某种语言后更换其编译器。 + +## 开发检查 + +clang-format、clang-tidy 和 IWYU 的范围、依赖及手动使用方式见 +[代码格式和静态检查](scripts/README.md)。 + +## 许可证 + +InfiniTrain Backends 使用 [MIT 许可证](LICENSE)。 diff --git a/backends/maca/CMakeLists.txt b/backends/maca/CMakeLists.txt new file mode 100644 index 0000000..613e497 --- /dev/null +++ b/backends/maca/CMakeLists.txt @@ -0,0 +1,102 @@ +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 + ${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. + block(SCOPE_FOR VARIABLES) + # InfiniTrain v0.7.0 resolves common test sources via PROJECT_SOURCE_DIR. + # Scope that path to InfiniTrain while creating its provider test targets. + set(PROJECT_SOURCE_DIR "${INFINITRAIN_SOURCE_DIR}") + 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 + ) + endblock() +endif() + +if(INFINITRAIN_BACKENDS_BUILD_EXAMPLES) + add_subdirectory(examples) +endif() diff --git a/backends/maca/README.md b/backends/maca/README.md new file mode 100644 index 0000000..79cc1c6 --- /dev/null +++ b/backends/maca/README.md @@ -0,0 +1,191 @@ +# MACA Backend + +**English** | [简体中文](README.zh-CN.md) + +This backend provides MACA runtime support, kernels, and optional MCCL collectives +for InfiniTrain's GPT-2, LLaMA 3, and Mixtral examples. + +See the [project README](../../README.md) for the project overview, shared build +options, and architecture. Unless otherwise noted, paths and commands below are +relative to the repository root. + +## Requirements + +In addition to the [common requirements](../../README.md#common-requirements): + +- A compatible **MACA SDK** with a C++20-capable `mxgpu_llvm/bin/mxcc` compiler, + the MACA runtime, MCDNN, MCBLAS, and the SDK's OpenMP runtime +- **MCCL** when distributed collectives are enabled (enabled by default) +- **GNU Make** for the quick-start build commands +- Python 3 when using the automated model test runner (standard library only) + +## 🚀 Quick Start + +Starting from the repository root, initialize the pinned InfiniTrain submodule, +select the MACA SDK, and build in a dedicated directory: + +```bash +git submodule update --init --recursive +export MACA_PATH=/opt/maca + +mkdir build +cd build +cmake .. -DINFINITRAIN_BACKEND=maca -DINFINITRAIN_MACA_WITH_MCCL=ON +make -j +``` + +The example executables are written to the current `build` directory. To view +their command-line options, run: + +```bash +./gpt2 --help +./llama3 --help +./mixtral --help +``` + +For instructions on running specific test cases, see the +[InfiniTrain README](https://github.com/InfiniTensor/InfiniTrain#readme). + +## MACA Build Options + +| Option | Default | Description | +| ------ | ------- | ----------- | +| `INFINITRAIN_MACA_WITH_MCCL` | `ON` | Enable MCCL distributed collectives | +| `MACA_PATH` | `$MACA_PATH` | MACA SDK root | + +`MACA_PATH` can be set through the environment or with `-DMACA_PATH=/path/to/maca`. +To build without MCCL, pass `-DINFINITRAIN_MACA_WITH_MCCL=OFF`. + +## 🏋️ Training + +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 /path/to/training_data.bin \ + --llmc_filepath /path/to/model.bin \ + --num_iteration 10 +``` + +GPT-2 and LLaMA 3 use `--batch_size` and `--total_batch_size`; Mixtral +uses `--micro_batch_size` and `--global_batch_size`. Run an executable with +`--help` to inspect its model and dataset options. + +Model tests cover combinations of models, training precision, and parallel +configurations. They reuse InfiniTrain's test runner and +`third_party/InfiniTrain/scripts/test_config.json`. The MACA entry point merges +that configuration with the overrides in +`backends/maca/scripts/test_config_maca.json`. All upstream groups and tests are +inherited in their original order. Entries in the overrides' `test_groups` only +customize settings; groups and tests omitted there are inherited as well. +Before running it, update the dataset and checkpoint paths in the overrides. +The wrapper uses Python's standard library for JSON handling, including the +upstream runner's configuration queries, so `jq` is not required. +The runner uses `build/maca` by default and clears that directory before each +rebuild and after the run. + +Run all configured model tests from the repository root in a separate shell +with the MACA SDK environment configured: + +```bash +backends/maca/scripts/run_models_and_profile.bash +``` + +Without `--only-run`, the runner executes all enabled test groups in the +configuration. To run just the basic group, add `--only-run basic`. + +## Unit Tests + +Unit tests are disabled by default. From the `build` directory created in the +quick start, enable and build them with: + +```bash +cmake .. -DBUILD_TEST=ON +make -j +ctest -L maca --output-on-failure +``` + +The generated MACA unit test binaries are named `test_*_maca`, contain only +`PRIVATEUSE1/*` GTest instances, and carry the `maca` CTest label. This label +identifies the backend and can be selected with CTest's `-L` option. + +The same build also contains InfiniTrain's CPU, fake-provider, and CPU-only +tests. Use `ctest -L cpu` for the upstream CPU tests, or run +`ctest --output-on-failure` without `-L` to execute the complete registered set. + +## Application Integration + +Applications must register the provider before parsing or constructing a +`maca` device: + +```cpp +#include "infini_train/include/device.h" +#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()` registers the process-wide PrivateUse1 name `maca` and +sets its default autocast dtype to BF16. Kernels, the device guard, and optional +MCCL support register through static initialization; the device runtime 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 uses `--whole-archive` to retain static registration +objects from InfiniTrain's core, CPU kernels, and the MACA provider. Library +targets that do not produce a final executable may link +`InfiniTrain::Backend::MACA`. + +The MACA implementation uses InfiniTrain's existing registration mechanisms: +`REGISTER_KERNEL`, `INFINI_TRAIN_REGISTER_DEVICE_GUARD_IMPL`, and +`INFINI_TRAIN_REGISTER_CCL_IMPL`. These registrations run during static +initialization; `RegisterBackend()` registers the provider metadata before +`maca` device parsing. + +When adding the MACA backend as a third-party library with `add_subdirectory()`, +select `${MACA_PATH}/mxgpu_llvm/bin/mxcc` in the parent project before enabling +C++. + +## Runtime Notes + +Configure these independent boolean flags in each model test's `args` in +`backends/maca/scripts/test_config_maca.json`, or pass them directly to the +example executable. The multithread workaround defaults to `true`, pool +retention defaults to `false`, and the runtime does not infer a mode from the +thread or process count. + +| Flag | Default | Behavior | +| ---- | ------- | -------- | +| `--maca_multithread_workarounds=true\|false` | `true` | Enables synchronous allocation/free, the copy mutex, default launch blocking, and the existing example synchronization/exit workarounds. | +| `--maca_retain_async_pool=true\|false` | `false` | Retains pages in the SDK default async pool for reuse by setting its release threshold to `UINT64_MAX`. | + +The JSON omits redundant `maca_multithread_workarounds=true` entries and +explicitly disables that flag for cases using the native async path: basic +cases `1`, `1_bfloat16`, `2`, `2_bfloat16`, and every `8_proc` case. It enables +pool retention for those same cases. Pool retention avoids the reproduced +MACA 3.5.3.18 ATU fault in mixed parallel training by keeping idle pages +available for reuse; explicit trimming can reintroduce the fault. + +Parse flags before first device use. The runtime captures their values once. +Explicit `MACA_LAUNCH_BLOCKING` and `MCCL_P2P_DISABLE` environment values retain +precedence. When workarounds are enabled, the provider defaults +`MACA_LAUNCH_BLOCKING` to `1` and also defaults `MCCL_P2P_DISABLE` to `1` when +the command line specifies `--tensor_parallel` greater than `1`. With +workarounds disabled it sets neither +variable; multiprocess TP therefore needs separate validation with the SDK's +native P2P behavior. Example workarounds stay in InfiniTrain and are selected +by one optional build-time flag binding. Large MCCL all-reduce messages are +segmented at 160 MiB by default; set `INFINI_MCCL_ALLREDUCE_SEGMENT_MB=0` to +disable segmentation. diff --git a/backends/maca/README.zh-CN.md b/backends/maca/README.zh-CN.md new file mode 100644 index 0000000..cda4275 --- /dev/null +++ b/backends/maca/README.zh-CN.md @@ -0,0 +1,170 @@ +# MACA 后端 + +[English](README.md) | **简体中文** + +本后端为 InfiniTrain 的 GPT-2、LLaMA 3 和 Mixtral 示例提供 MACA 运行时支持、 +算子和可选的 MCCL 集合通信。 + +项目介绍、通用构建选项和架构见[项目 README](../../README.zh-CN.md)。 +除非另有说明，以下路径和命令均以仓库根目录为基准。 + +## 依赖 + +除[通用依赖](../../README.zh-CN.md#通用依赖)外，还需要： + +- 兼容的 **MACA SDK**，包含支持 C++20 的 `mxgpu_llvm/bin/mxcc` 编译器、MACA + 运行时、MCDNN、MCBLAS 及 SDK 自带的 OpenMP 运行时 +- 启用分布式集合通信时需要 **MCCL**（默认启用） +- 快速开始中的构建命令需要 **GNU Make** +- 使用自动化模型测试脚本时需要 Python 3（仅使用标准库） + +## 🚀 快速开始 + +从仓库根目录执行以下命令，初始化锁定版本的 InfiniTrain 子模块、选择 MACA SDK， +并在独立目录中构建： + +```bash +git submodule update --init --recursive +export MACA_PATH=/opt/maca + +mkdir build +cd build +cmake .. -DINFINITRAIN_BACKEND=maca -DINFINITRAIN_MACA_WITH_MCCL=ON +make -j +``` + +示例可执行文件输出到当前的 `build` 目录。可直接查看各示例的命令行选项： + +```bash +./gpt2 --help +./llama3 --help +./mixtral --help +``` + +具体测例的运行方式请参考 [InfiniTrain README](https://github.com/InfiniTensor/InfiniTrain#readme)。 + +## MACA 构建选项 + +| 选项 | 默认值 | 说明 | +| ---- | ------ | ---- | +| `INFINITRAIN_MACA_WITH_MCCL` | `ON` | 启用 MCCL 分布式集合通信 | +| `MACA_PATH` | `$MACA_PATH` | MACA SDK 根目录 | + +`MACA_PATH` 可通过环境变量或 `-DMACA_PATH=/path/to/maca` 设置。 +如需在不使用 MCCL 的情况下构建，可传入 `-DINFINITRAIN_MACA_WITH_MCCL=OFF`。 + +## 🏋️ 训练 + +各模型示例都是独立的可执行文件。通过 `--device maca` +选择 MACA 后端。例如，在构建目录中执行以下命令可启动单节点 LLaMA 3 训练： + +```bash +./llama3 \ + --device maca \ + --input_bin /path/to/training_data.bin \ + --llmc_filepath /path/to/model.bin \ + --num_iteration 10 +``` + +GPT-2 和 LLaMA 3 使用 `--batch_size` 与 `--total_batch_size`；Mixtral 使用 +`--micro_batch_size` 与 `--global_batch_size`。可通过各可执行文件的 `--help` +查看模型及数据集相关选项。 + +模型测试覆盖不同模型、训练精度和并行配置的组合，复用 InfiniTrain 的测试脚本和 +`third_party/InfiniTrain/scripts/test_config.json`。MACA 入口脚本将该配置与 +`backends/maca/scripts/test_config_maca.json` 中的覆盖配置合并，继承上游全部测试组 +及组内测例，并保留上游顺序。覆盖配置中的 `test_groups` 仅用于修改参数；未列出的组 +和测例也会继承。运行前，请修改 +覆盖配置中的数据集和 checkpoint 路径。 +入口脚本使用 Python 标准库处理 JSON，包括上游脚本的配置查询，无需安装 `jq`。 +自动测试默认在 `build/maca` 中构建，脚本会在每次重新构建前和运行结束后清理该目录。 + +请在已配置 MACA SDK 环境的另一个 shell 中，从仓库根目录运行全部已配置的模型 +测试： + +```bash +backends/maca/scripts/run_models_and_profile.bash +``` + +不传 `--only-run` 时，脚本会运行配置中启用的全部测试组。如只需运行 basic 组， +可添加 `--only-run basic`。 + +## 单元测试 + +单元测试默认关闭。在快速开始创建的 `build` 目录中，执行以下命令启用并构建 +单元测试： + +```bash +cmake .. -DBUILD_TEST=ON +make -j +ctest -L maca --output-on-failure +``` + +生成的 MACA 单元测试程序命名为 `test_*_maca`，仅包含 `PRIVATEUSE1/*` GTest +实例，并带有 `maca` CTest 标签。该标签标识后端，可通过 CTest 的 `-L` 选项筛选。 + +同一构建中还包含 InfiniTrain 的 CPU 测试、模拟后端测试和仅适用于 CPU 的测试。 +使用 `ctest -L cpu` 可运行上游 CPU 测试；使用不带 `-L` 的 +`ctest --output-on-failure` 可运行全部已注册测试。 + +## 应用集成 + +应用应在解析或构造 `maca` 设备前注册后端： + +```cpp +#include "infini_train/include/device.h" +#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()` 注册进程级 PrivateUse1 名称 `maca`，并将其默认 autocast +数据类型设为 BF16。算子、设备 guard 和可选的 MCCL 支持通过静态初始化注册； +设备运行时延迟到 InfiniTrain 首次请求设备 guard 时才初始化。 + +最终可执行文件应链接 `InfiniTrain::Backend::MACAExecutable`，由该目标处理静态库 +的链接： + +```cmake +add_executable(train main.cc) +target_link_libraries(train PRIVATE InfiniTrain::Backend::MACAExecutable) +``` + +该接口目标使用 `--whole-archive` 保留 InfiniTrain 核心、CPU 算子和 MACA 后端的 +静态注册对象。不生成最终可执行文件的库目标可链接 `InfiniTrain::Backend::MACA`。 + +MACA 实现使用 InfiniTrain 现有的注册机制：`REGISTER_KERNEL`、 +`INFINI_TRAIN_REGISTER_DEVICE_GUARD_IMPL` 和 `INFINI_TRAIN_REGISTER_CCL_IMPL`。 +这些注册在静态初始化阶段执行；`RegisterBackend()` 在解析 `maca` 设备前注册 +后端元数据。 + +通过 `add_subdirectory()` 将 MACA 后端作为第三方库添加到父项目时，应在启用 C++ +前选择 `${MACA_PATH}/mxgpu_llvm/bin/mxcc` 编译器。 + +## 运行时说明 + +以下两个布尔参数相互独立，可在 `backends/maca/scripts/test_config_maca.json` +中各模型测试的 `args` 中配置，也可直接传给示例可执行文件。多线程兼容措施默认为 +`true`，异步内存池页面保留默认为 `false`，运行时不会根据线程数或进程数推断模式。 + +| 参数 | 默认值 | 行为 | +| ---- | ------ | ---- | +| `--maca_multithread_workarounds=true\|false` | `true` | 启用同步分配与释放、内存复制互斥锁、默认的阻塞式 launch，以及现有示例中的同步和退出兼容措施。 | +| `--maca_retain_async_pool=true\|false` | `false` | 将 SDK 默认异步内存池的释放阈值设为 `UINT64_MAX`，保留页面以便复用。 | + +JSON 配置省略了冗余的 `maca_multithread_workarounds=true`，并对使用原生异步 +路径的用例显式关闭该参数：包括 basic 组中的 `1`、`1_bfloat16`、`2`、 +`2_bfloat16`，以及所有 `8_proc` 用例。这些用例同时启用了内存池页面保留。 +在混合并行训练中，保留空闲页面供后续复用，可避开已复现的 MACA 3.5.3.18 ATU +故障；显式裁剪内存池可能再次触发该故障。 + +请在首次使用设备前解析参数，运行时只捕获一次参数值。显式设置的 +`MACA_LAUNCH_BLOCKING` 和 `MCCL_P2P_DISABLE` 环境变量优先。在启用兼容措施时， +后端默认将 `MACA_LAUNCH_BLOCKING` 设为 `1`；如果命令行指定的 +`--tensor_parallel` 大于 `1`，还默认将 `MCCL_P2P_DISABLE` 设为 `1`。关闭兼容 +措施时，后端不会设置这两个变量，因此多进程 TP 仍需单独验证 SDK 的原生 P2P +行为。示例中的兼容措施保留在 InfiniTrain 中，通过一个可选的构建期参数绑定来 +控制。较大的 MCCL all-reduce 消息默认按 160 MiB 分段；设置 +`INFINI_MCCL_ALLREDUCE_SEGMENT_MB=0` 可禁用分段。 diff --git a/backends/maca/cmake/pre_project.cmake b/backends/maca/cmake/pre_project.cmake new file mode 100644 index 0000000..a6e6a0e --- /dev/null +++ b/backends/maca/cmake/pre_project.cmake @@ -0,0 +1,85 @@ +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; +# in particular, glog's PIE probes use -fPIE, which mxcc does not support. +# 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_ELF_H 1) +set(HAVE_GLOB_H 1) +set(HAVE_LINK_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..26b2625 --- /dev/null +++ b/backends/maca/examples/CMakeLists.txt @@ -0,0 +1,38 @@ +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 + INFINITRAIN_EXAMPLE_EXTERNAL_BACKEND_WORKAROUNDS=FLAGS_maca_multithread_workarounds + ) + 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) +add_maca_model_example(qwen3) 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..7b7fa86 --- /dev/null +++ b/backends/maca/include/infini_train_maca/backend.h @@ -0,0 +1,14 @@ +#pragma once + +#include "gflags/gflags_declare.h" + +DECLARE_bool(maca_multithread_workarounds); +DECLARE_bool(maca_retain_async_pool); + +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..d9af67b --- /dev/null +++ b/backends/maca/scripts/run_models_and_profile.bash @@ -0,0 +1,174 @@ +#!/usr/bin/env bash + +set -euo pipefail + +SCRIPT_DIR="$(cd -- "$(dirname -- "${BASH_SOURCE[0]}")" && pwd)" +REPOSITORY_ROOT="$(cd -- "${SCRIPT_DIR}/../../.." && pwd)" +MACA_CONFIG="${SCRIPT_DIR}/test_config_maca.json" +BASE_CONFIG="${REPOSITORY_ROOT}/third_party/InfiniTrain/scripts/test_config.json" + +if ! command -v python3 >/dev/null 2>&1; then + echo "Error: python3 is required by the model test runner." >&2 + exit 1 +fi + +# Reuse the Python standard library already used by upstream's log tools. +# Keep configuration handling here so no helper file or Python package is needed. +json_config() { + python3 - "$@" <<'PY' +import json +import re +import sys + + +def load(path): + with open(path, encoding="utf-8") as stream: + return json.load(stream) + + +def merge(base, overrides): + result = base.copy() + for key, value in overrides.items(): + if isinstance(result.get(key), dict) and isinstance(value, dict): + result[key] = merge(result[key], value) + else: + result[key] = value + return result + + +def text(value): + if isinstance(value, str): + return value + return json.dumps(value, ensure_ascii=False, separators=(",", ":")) + + +def namespaced_path(path, model, mode): + match = re.search(r"^(.*)/(checkpoint_step_[0-9]+(?:/.*)?)$", path) + if match: + return f"{match[1]}/{model}/{mode}/{match[2]}" + return f"{path}/{model}/{mode}" + + +def query(arguments): + # This adapter supports the pinned runner's queries, not arbitrary jq code. + values = {} + while arguments and arguments[0].startswith("-"): + option = arguments.pop(0) + if option == "-r": + continue + if option not in ("--arg", "--argjson"): + raise ValueError(f"Unsupported upstream jq option: {option}") + key, value = arguments[:2] + del arguments[:2] + values[key] = json.loads(value) if option == "--argjson" else value + expression, path = arguments + config = load(path) + expression = " ".join(expression.split()) + + if expression == '.variables[$k] // empty': + value = config["variables"].get(values["k"]) + if value is not None and value is not False: + print(text(value)) + return + if expression == r'.variables | to_entries[] | "\(.key)=\(.value)"': + for key, value in config["variables"].items(): + print(f"{key}={text(value)}") + return + if expression.startswith("def namespaced_path("): + args = config["test_groups"][values["g"]]["tests"][values["t"]]["args"].copy() + model = values["model"] + mode = "resume" if "load" in args else "no_resume" + source_mode = "no_resume" if "no_resume" in (args.get("load") or "") else "resume" + if "save" in args: + args["save"] = namespaced_path(args["save"], model, mode) + if "load" in args: + args["load"] = namespaced_path(args["load"], model, source_mode) + for key, value in args.items(): + print(f"--{key}={text(value)}") + return + + match = re.fullmatch(r"((?:\.[A-Za-z_]\w*|\[\d+\])+)( \| length| // empty)?", expression) + if not match: + raise ValueError(f"Unsupported upstream jq query: {expression}") + value = config + for key, index in re.findall(r"\.([A-Za-z_]\w*)|\[(\d+)\]", match[1]): + if value is None: + break + value = value.get(key) if key else value[int(index)] + if match[2] == " // empty" and (value is None or value is False): + return + print(text(len(value) if match[2] == " | length" else value)) + + +def main(): + action, *arguments = sys.argv[1:] + if action == "source-dir": + for command in load(arguments[0])["basic_compile_commands"]: + match = re.search(r"-DINFINITRAIN_SOURCE_DIR=([^ ]+)", command["cmd"]) + if match: + print(match[1]) + break + elif action == "merge": + base, overrides = map(load, arguments) + config = merge(base, {k: v for k, v in overrides.items() if k != "test_groups"}) + group_overrides = {group["tag"]: group for group in overrides.get("test_groups", [])} + upstream_tags = {group["tag"] for group in base["test_groups"]} + for tag in group_overrides: + if tag not in upstream_tags: + raise ValueError(f"Unknown upstream test group: {tag}") + + # Inherit every upstream group; MACA entries only override its settings. + config["test_groups"] = [] + for group in base["test_groups"]: + tag = group["tag"] + override = group_overrides.get(tag, {}) + test_ids = {test["id"] for test in group["tests"]} + unknown_ids = [test["id"] for test in override.get("tests", []) if test["id"] not in test_ids] + if unknown_ids: + raise ValueError(f"Unknown upstream tests in {tag}: {', '.join(unknown_ids)}") + merged_group = merge(group, {k: v for k, v in override.items() if k != "tests"}) + merged_group["tests"] = [ + merge(test, next((item for item in override.get("tests", []) if item["id"] == test["id"]), {})) + for test in group["tests"] + ] + config["test_groups"].append(merged_group) + json.dump(config, sys.stdout, ensure_ascii=False, indent=2) + print() + elif action == "query": + query(arguments) + else: + raise ValueError(f"Unknown configuration action: {action}") + + +try: + main() +except (OSError, ValueError, KeyError, IndexError, TypeError) as error: + sys.exit(f"Error: {error}") +PY +} + +CONFIGURED_INFINITRAIN_SOURCE_DIR="$(json_config source-dir "${MACA_CONFIG}")" +INFINITRAIN_SOURCE_DIR="${INFINITRAIN_SOURCE_DIR:-${CONFIGURED_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 merge MACA overrides into the +# pinned upstream configuration before passing it to the runner. +cd "${REPOSITORY_ROOT}" +MERGED_CONFIG="$(mktemp "${TMPDIR:-/tmp}/infinitrain-maca-config.XXXXXX.json")" +trap 'rm -f -- "${MERGED_CONFIG}"' EXIT +json_config merge "${BASE_CONFIG}" "${MACA_CONFIG}" > "${MERGED_CONFIG}" + +# Keep the adapter local to the runner; build commands inherit no jq replacement. +# Upstream uses set -e/pipefail but does not enable nounset. +( + set +u + jq() { json_config query "$@"; } + source "${UPSTREAM_RUNNER}" --test-config "${MERGED_CONFIG}" "$@" +) diff --git a/backends/maca/scripts/test_config_maca.json b/backends/maca/scripts/test_config_maca.json new file mode 100644 index 0000000..22a5c48 --- /dev/null +++ b/backends/maca/scripts/test_config_maca.json @@ -0,0 +1,139 @@ +{ + "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", + "LLAMA3_LORA_WEIGHTS_DIR": "/nfs/InfiniTrain-dev/data/llmc/llama3/llama3.2_1B_lora_weights_rank8_alpha16.bin", + "QWEN3_INPUT_BIN": "/nfs/InfiniTrain-dev/data/llmc/qwen3/tinyshakespeare/tiny_shakespeare_train.bin", + "QWEN3_LLMC_FILEPATH": "/nfs/InfiniTrain-dev/data/llmc/qwen3/qwen3-8b-fp32.llmc", + "QWEN3_LORA_WEIGHTS_DIR": "/nfs/InfiniTrain-dev/data/llmc/qwen3/qwen3_8b_lora_weights_rank8_alpha16.bin", + "GPT2_LORA_WEIGHTS_DIR": "/nfs/InfiniTrain-dev/data/llmc/gpt2/gpt2_124M_lora_weights_rank8_alpha16.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" + }, + "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": { + "maca_multithread_workarounds": false, + "maca_retain_async_pool": true + } + }, + { + "id": "1_bfloat16", + "args": { + "maca_multithread_workarounds": false, + "maca_retain_async_pool": true + } + }, + { + "id": "2", + "args": { + "maca_multithread_workarounds": false, + "maca_retain_async_pool": true + } + }, + { + "id": "2_bfloat16", + "args": { + "maca_multithread_workarounds": false, + "maca_retain_async_pool": true + } + } + ] + }, + { + "tag": "8_proc", + "tests": [ + { + "id": "3_8proc", + "args": { + "maca_multithread_workarounds": false, + "maca_retain_async_pool": true + } + }, + { + "id": "3_bfloat16_8proc", + "args": { + "maca_multithread_workarounds": false, + "maca_retain_async_pool": true + } + }, + { + "id": "4_8proc", + "args": { + "maca_multithread_workarounds": false, + "maca_retain_async_pool": true + } + }, + { + "id": "4_bfloat16_8proc", + "args": { + "maca_multithread_workarounds": false, + "maca_retain_async_pool": true + } + }, + { + "id": "5_8proc", + "args": { + "maca_multithread_workarounds": false, + "maca_retain_async_pool": true + } + }, + { + "id": "5_bfloat16_8proc", + "args": { + "maca_multithread_workarounds": false, + "maca_retain_async_pool": true + } + }, + { + "id": "6_8proc", + "args": { + "maca_multithread_workarounds": false, + "maca_retain_async_pool": true + } + }, + { + "id": "6_bfloat16_8proc", + "args": { + "maca_multithread_workarounds": false, + "maca_retain_async_pool": true + } + }, + { + "id": "8_8proc", + "args": { + "maca_multithread_workarounds": false, + "maca_retain_async_pool": true + } + }, + { + "id": "8_bfloat16_8proc", + "args": { + "maca_multithread_workarounds": false, + "maca_retain_async_pool": true + } + } + ] + } + ] +} diff --git a/backends/maca/src/backend.cc b/backends/maca/src/backend.cc new file mode 100644 index 0000000..ebac53d --- /dev/null +++ b/backends/maca/src/backend.cc @@ -0,0 +1,21 @@ +#include "infini_train_maca/backend.h" + +#include + +#include "infini_train/include/core/privateuse1_backend.h" + +#include "runtime/maca_guard_impl.h" + +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; + 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..4d37e28 --- /dev/null +++ b/backends/maca/src/ccl/mccl_impl.cc @@ -0,0 +1,201 @@ +#include "ccl/mccl_impl.h" + +#include +#include +#include +#include +#include +#include +#include +#include + +#include + +#include "glog/logging.h" + +#include "infini_train/include/core/runtime/runtime_common.h" +#include "infini_train/include/device.h" +#include "infini_train/include/nn/parallel/reduce_op_type.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::comm::ReduceOpType::kSum, mcclSum}, {nn::parallel::comm::ReduceOpType::kProd, mcclProd}, + {nn::parallel::comm::ReduceOpType::kMin, mcclMin}, {nn::parallel::comm::ReduceOpType::kMax, mcclMax}, + {nn::parallel::comm::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::comm::ReduceOpType reduce_op, const CclComm *comm, Stream *stream) const { + // C550 all-reduce bandwidth drops sharply for very large messages. Segment + // at the provider boundary so every framework caller stays below the cliff. + constexpr size_t kBytesPerMB = 1024ULL * 1024ULL; + constexpr size_t kDefaultSegmentBytes = 160 * kBytesPerMB; + static const size_t segment_bytes = [] { + const char *value = std::getenv("INFINI_MCCL_ALLREDUCE_SEGMENT_MB"); + if (value == nullptr || *value == '\0') { + return kDefaultSegmentBytes; + } + const std::string_view text(value); + size_t size_mb = 0; + const auto [end, error] = std::from_chars(text.data(), text.data() + text.size(), size_mb); + CHECK(error == std::errc{} && end == text.data() + text.size() && size_mb <= 4096) + << "INFINI_MCCL_ALLREDUCE_SEGMENT_MB must be an integer in [0, 4096]"; + return size_mb == 0 ? std::numeric_limits::max() : size_mb * kBytesPerMB; + }(); + + const size_t element_size = infini_train::kDataTypeToSize.at(dtype); + const size_t segment_elements = std::max(1, segment_bytes / element_size); + const auto mccl_dtype = kMcclDtypeMap.at(dtype); + const auto mccl_op = kMcclReduceOpMap.at(reduce_op); + auto mccl_comm = GetMcclComm(comm); + auto mccl_stream = GetMacaStream(stream); + if (count <= segment_elements) { + MCCL_CHECK(mcclAllReduce(sendbuff, recvbuff, count, mccl_dtype, mccl_op, mccl_comm, mccl_stream)); + return; + } + + const auto *bytes = static_cast(sendbuff); + auto *output = static_cast(recvbuff); + for (size_t offset = 0; offset < count;) { + const size_t elements = std::min(segment_elements, count - offset); + MCCL_CHECK(mcclAllReduce(bytes + offset * element_size, output + offset * element_size, elements, mccl_dtype, + mccl_op, mccl_comm, mccl_stream)); + offset += elements; + } +} + +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::comm::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::comm::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))); +} + +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..4af2056 --- /dev/null +++ b/backends/maca/src/ccl/mccl_impl.h @@ -0,0 +1,49 @@ +#pragma once + +#include + +#include "infini_train/include/core/ccl/ccl.h" + +namespace infini_train::core::maca { + +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::comm::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::comm::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::comm::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..9cc6c91 --- /dev/null +++ b/backends/maca/src/kernels/accumulate_grad.maca @@ -0,0 +1,98 @@ +#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/kernel_helper.cuh" +#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 { + +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..4d74689 --- /dev/null +++ b/backends/maca/src/kernels/cast.maca @@ -0,0 +1,58 @@ +#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 { + +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..a0cbe67 --- /dev/null +++ b/backends/maca/src/kernels/common/gemm.cc @@ -0,0 +1,97 @@ +#include "kernels/common/gemm.h" + +#include +#include + +#include + +#include "glog/logging.h" + +#include "infini_train/include/core/runtime/device_guard.h" +#include "infini_train/include/dispatcher.h" + +#include "common/common_maca.h" +#include "runtime/maca_runtime_common.h" + +namespace infini_train::kernels::maca { +namespace { + +// NOTE(cx): Developer switch: set INFINI_MACA_GEMM_FP64_ACCUM=1 before launch for FP64 +// multiply-accumulate on FP32 GEMMs, then round the output to FP32. Default off; +// cached on first GEMM. Uses extra FP64 buffers and compute; BF16 is unaffected. +bool Fp64AccumEnabled() { + static const bool enabled = [] { + const char *value = std::getenv("INFINI_MACA_GEMM_FP64_ACCUM"); + if (value == nullptr || *value == '\0' || std::string_view(value) == "0") { + return false; + } + CHECK(std::string_view(value) == "1") << "INFINI_MACA_GEMM_FP64_ACCUM must be 0 or 1"; + LOG(INFO) << "MACA GEMM FP64 accumulation enabled for FP32 inputs/output (developer mode)"; + return true; + }(); + return enabled; +} + +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); + + if (Fp64AccumEnabled() && params.input_dtype == DataType::kFLOAT32 && params.output_dtype == DataType::kFLOAT32) { + GemmFp64Accum(device, params, blas_handle); + return; + } + + 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)); + } +} + +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..bae56c1 --- /dev/null +++ b/backends/maca/src/kernels/common/gemm.h @@ -0,0 +1,14 @@ +#pragma once + +#include + +#include "infini_train/src/kernels/common/gemm.h" + +namespace infini_train::kernels::maca { + +void Gemm(Device device, GemmParams params); + +// Developer reference path: FP32 inputs/output with FP64 GEMM temporaries. +void GemmFp64Accum(Device device, GemmParams params, mcblasHandle_t blas_handle); + +} // 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..f382120 --- /dev/null +++ b/backends/maca/src/kernels/concat.maca @@ -0,0 +1,253 @@ +#include +#include +#include +#include +#include +#include +#include +#include + +#include "glog/logging.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 "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 { + +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..51f1c7e --- /dev/null +++ b/backends/maca/src/kernels/cross_entropy.maca @@ -0,0 +1,237 @@ +#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): Replace the host FP64 accumulation below with a GPU FP32 hierarchical reduction to remove + // the D2H/H2D transfers and host synchronization. + 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()); + const double loss_sum + = std::accumulate(loss_cpu_typed_ptr, loss_cpu_typed_ptr + bs, 0.0, [](double acc, const Tinput &val) { + return acc + static_cast(common::maca::Cast(val)); + }); + const float loss_mean = static_cast(loss_sum / static_cast(bs)); + static_cast(loss->DataPtr())[0] = loss_mean; + + 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 { + +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..2b89f36 --- /dev/null +++ b/backends/maca/src/kernels/elementwise.maca @@ -0,0 +1,1296 @@ +#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 { + +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..922f29d --- /dev/null +++ b/backends/maca/src/kernels/embedding.maca @@ -0,0 +1,134 @@ +#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 "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 { + +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..db72d91 --- /dev/null +++ b/backends/maca/src/kernels/fill.maca @@ -0,0 +1,57 @@ +#include +#include +#include + +#include "infini_train/include/core/runtime/device_guard.h" +#include "infini_train/include/device.h" +#include "infini_train/include/dispatcher.h" +#include "infini_train/include/tensor.h" + +#include "common/kernel_helper.cuh" +#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 { + +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..9caa710 --- /dev/null +++ b/backends/maca/src/kernels/gather.maca @@ -0,0 +1,235 @@ +#include +#include +#include +#include +#include +#include + +#include "glog/logging.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 "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 { + +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..ecd8942 --- /dev/null +++ b/backends/maca/src/kernels/layernorm.maca @@ -0,0 +1,246 @@ +#include +#include +#include +#include +#include + +#include + +#include "infini_train/include/core/runtime/device_guard.h" +#include "infini_train/include/device.h" +#include "infini_train/include/dispatcher.h" +#include "infini_train/include/tensor.h" + +#include "common/common_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 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 { + +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..f3536e2 --- /dev/null +++ b/backends/maca/src/kernels/linear.maca @@ -0,0 +1,580 @@ +#include +#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 { +namespace { + +constexpr int kThreads = 256; + +int Blocks(int64_t count) { return static_cast(std::min((count + kThreads - 1) / kThreads, 65535)); } + +// Pack only matrix elements, preserving the caller's padding and broadcast strides. +template +__global__ void CopyMatrix(const Src *src, Dst *dst, int rows, int cols, int src_ld, int dst_ld, int64_t src_stride, + int64_t dst_stride, int64_t count) { + const int64_t matrix_size = static_cast(rows) * cols; + for (int64_t idx = static_cast(blockIdx.x) * blockDim.x + threadIdx.x; idx < count; + idx += static_cast(gridDim.x) * blockDim.x) { + const int64_t batch = idx / matrix_size; + const int64_t element = idx % matrix_size; + const int64_t row = element % rows; + const int64_t col = element / rows; + dst[batch * dst_stride + col * dst_ld + row] = static_cast(src[batch * src_stride + col * src_ld + row]); + } +} + +__global__ void ScaleOutput(float *output, int rows, int cols, int ld, int64_t stride, int64_t count, double beta) { + const int64_t matrix_size = static_cast(rows) * cols; + for (int64_t idx = static_cast(blockIdx.x) * blockDim.x + threadIdx.x; idx < count; + idx += static_cast(gridDim.x) * blockDim.x) { + const int64_t batch = idx / matrix_size; + const int64_t element = idx % matrix_size; + const int64_t offset = batch * stride + (element / rows) * ld + element % rows; + output[offset] = beta == 0.0 ? 0.0f : static_cast(beta * static_cast(output[offset])); + } +} + +} // namespace + +void GemmFp64Accum(Device device, GemmParams params, mcblasHandle_t blas_handle) { + CHECK_GE(params.m, 0); + CHECK_GE(params.n, 0); + CHECK_GE(params.k, 0); + CHECK_GE(params.batch_count, 0); + if (params.m == 0 || params.n == 0 || params.batch_count == 0) { + return; + } + + mcStream_t stream = nullptr; + MCBLAS_CHECK(mcblasGetStream(blas_handle, &stream)); + const int64_t c_size = static_cast(params.m) * params.n; + const int64_t c_count = c_size * params.batch_count; + const double alpha = params.alpha; + const double beta = params.beta; + if (params.k == 0 || alpha == 0.0) { + ScaleOutput<<>>(static_cast(params.C), params.m, params.n, + params.ldc, params.stride_c, c_count, beta); + MACA_CHECK(mcGetLastError()); + return; + } + + const bool transpose_a = params.trans_a == GemmTranspose::kTranspose; + const bool transpose_b = params.trans_b == GemmTranspose::kTranspose; + const int a_rows = transpose_a ? params.k : params.m; + const int a_cols = transpose_a ? params.m : params.k; + const int b_rows = transpose_b ? params.n : params.k; + const int b_cols = transpose_b ? params.k : params.n; + const int64_t a_size = static_cast(a_rows) * a_cols; + const int64_t b_size = static_cast(b_rows) * b_cols; + const int64_t a_count = a_size * params.batch_count; + const int64_t b_count = b_size * params.batch_count; + + // mcBLAS does not support FP32 A/B/C with COMPUTE_64F. Convert exact FP32 + // values to double, perform the full GEMM in double, then round C once. + Tensor a64({params.batch_count, a_cols, a_rows}, DataType::kFLOAT64, device); + Tensor b64({params.batch_count, b_cols, b_rows}, DataType::kFLOAT64, device); + Tensor c64({params.batch_count, params.n, params.m}, DataType::kFLOAT64, device); + auto *a = static_cast(a64.DataPtr()); + auto *b = static_cast(b64.DataPtr()); + auto *c = static_cast(c64.DataPtr()); + + CopyMatrix<<>>(static_cast(params.A), a, a_rows, a_cols, + params.lda, a_rows, params.stride_a, a_size, a_count); + MACA_CHECK(mcGetLastError()); + CopyMatrix<<>>(static_cast(params.B), b, b_rows, b_cols, + params.ldb, b_rows, params.stride_b, b_size, b_count); + MACA_CHECK(mcGetLastError()); + if (beta != 0.0) { + CopyMatrix<<>>(static_cast(params.C), c, params.m, + params.n, params.ldc, params.m, params.stride_c, c_size, + c_count); + MACA_CHECK(mcGetLastError()); + } + + const auto trans_a = transpose_a ? MCBLAS_OP_T : MCBLAS_OP_N; + const auto trans_b = transpose_b ? MCBLAS_OP_T : MCBLAS_OP_N; + if (params.batch_count == 1) { + MCBLAS_CHECK(mcblasDgemm(blas_handle, trans_a, trans_b, params.m, params.n, params.k, &alpha, a, a_rows, b, + b_rows, &beta, c, params.m)); + } else { + MCBLAS_CHECK(mcblasDgemmStridedBatched(blas_handle, trans_a, trans_b, params.m, params.n, params.k, &alpha, a, + a_rows, a_size, b, b_rows, b_size, &beta, c, params.m, c_size, + params.batch_count)); + } + + CopyMatrix<<>>(c, static_cast(params.C), params.m, params.n, + params.m, params.ldc, c_size, params.stride_c, c_count); + MACA_CHECK(mcGetLastError()); +} + +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]; +} + +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()); + + // d_input = d_output * weight (transposed weight layout) or d_output * weight.T. + auto device = grad_output->GetDevice(); + 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, + }); + + 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 { + +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..156a207 --- /dev/null +++ b/backends/maca/src/kernels/no_op.maca @@ -0,0 +1,35 @@ +#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 { + +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..3dfdeaa --- /dev/null +++ b/backends/maca/src/kernels/outer.maca @@ -0,0 +1,192 @@ +#include +#include +#include +#include + +#include + +#include "glog/logging.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 "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 { + +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..86858d6 --- /dev/null +++ b/backends/maca/src/kernels/reduction.maca @@ -0,0 +1,251 @@ +#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 { + +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/scatter.maca b/backends/maca/src/kernels/scatter.maca new file mode 100644 index 0000000..9504d1f --- /dev/null +++ b/backends/maca/src/kernels/scatter.maca @@ -0,0 +1,124 @@ +#include +#include +#include +#include +#include + +#include "glog/logging.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 "common/common_maca.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; +} + +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..8cd0bf6 --- /dev/null +++ b/backends/maca/src/kernels/slice.maca @@ -0,0 +1,208 @@ +#include +#include +#include + +#include "glog/logging.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 "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 { + +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..6ece1bf --- /dev/null +++ b/backends/maca/src/kernels/softmax.maca @@ -0,0 +1,226 @@ +#include +#include +#include +#include + +#include + +#include "glog/logging.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/cub_compat.cuh" +#include "common/kernel_helper.cuh" +#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 { + +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..9452978 --- /dev/null +++ b/backends/maca/src/kernels/split.maca @@ -0,0 +1,187 @@ +#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 "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 { + +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..80688a9 --- /dev/null +++ b/backends/maca/src/kernels/stack.maca @@ -0,0 +1,162 @@ +#include +#include +#include +#include +#include +#include +#include + +#include "glog/logging.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 "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 { + +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/swiglu.maca b/backends/maca/src/kernels/swiglu.maca new file mode 100644 index 0000000..4ac5202 --- /dev/null +++ b/backends/maca/src/kernels/swiglu.maca @@ -0,0 +1,149 @@ +// Adapted from InfiniTrain v0.7.0's CUDA SwiGLU kernels for packed [gate, up] inputs. +#include +#include +#include +#include + +#include "common/common_maca.h" +#include "common/kernel_helper.cuh" +#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/dispatcher.h" +#include "infini_train/include/tensor.h" + +#include "runtime/maca_runtime_common.h" + +namespace infini_train::kernels::maca { +namespace { +using namespace infini_train::common::maca; + +// TODO(zbl): Optimize the packed [gate, up] accesses with vectorized +// loads/stores. +template +__global__ void SwiGLUForwardKernel(T *__restrict__ output, const T *__restrict__ input, int64_t hidden, + 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) { + const size_t row = idx / hidden; + const size_t col = idx % hidden; + const size_t input_base = row * 2 * hidden; + const T gate = input[input_base + col]; + const T up = input[input_base + hidden + col]; + output[idx] = Mul(up, Mul(gate, Sigmoid(gate))); + } +} + +template +__global__ void SwiGLUBackwardKernel(T *__restrict__ grad_input, const InputT *__restrict__ input, + const GradT *__restrict__ grad_output, int64_t hidden, 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) { + const size_t row = idx / hidden; + const size_t col = idx % hidden; + const size_t input_base = row * 2 * hidden; + const T gate = Cast(input[input_base + col]); + const T up = Cast(input[input_base + hidden + col]); + const T grad = Cast(grad_output[idx]); + const T sigmoid = Sigmoid(gate); + grad_input[input_base + col] = Mul(grad, Mul(up, Mul(sigmoid, Add(T(1), Mul(gate, Sub(T(1), sigmoid)))))); + grad_input[input_base + hidden + col] = Mul(grad, Mul(gate, sigmoid)); + } +} + +} // namespace + +std::shared_ptr SwiGLUForward(const std::shared_ptr &input) { + CHECK(input->IsContiguous()); + auto output_dims = input->Dims(); + CHECK_GT(output_dims.size(), 0); + CHECK_EQ(output_dims.back() % 2, 0); + const int64_t hidden = output_dims.back() / 2; + CHECK_GT(hidden, 0); + output_dims.back() = hidden; + + auto output = std::make_shared(output_dims, input->Dtype(), input->GetDevice()); + auto device = output->GetDevice(); + const auto &stream = dynamic_cast( + infini_train::core::GetDeviceGuardImpl(device.type())->GetStream(device)) + ->maca_stream(); + const size_t num_elements = output->NumElements(); + if (num_elements == 0) { + return output; + } + const dim3 block(256); + const dim3 grid(std::min(CEIL_DIV(num_elements, block.x), static_cast(65535))); + + switch (input->Dtype()) { + case DataType::kFLOAT32: + SwiGLUForwardKernel<<>>(static_cast(output->DataPtr()), + static_cast(input->DataPtr()), hidden, + num_elements); + break; + case DataType::kBFLOAT16: + SwiGLUForwardKernel<<>>(static_cast<__maca_bfloat16 *>(output->DataPtr()), + static_cast(input->DataPtr()), hidden, + num_elements); + break; + default: + LOG_LOC(FATAL, "MACA SwiGLUForward: unsupported data type"); + } + return output; +} + +std::shared_ptr SwiGLUBackward(const std::shared_ptr &input, + const std::shared_ptr &grad_output) { + CHECK(input->IsContiguous()); + CHECK(grad_output->IsContiguous()); + CHECK_GT(input->Dims().size(), 0); + const int64_t hidden = input->Dims().back() / 2; + CHECK_GT(hidden, 0); + CHECK_EQ(grad_output->NumElements() * 2, input->NumElements()); + + const DataType output_dtype = PromoteDataTypes(input->Dtype(), grad_output->Dtype()); + auto grad_input = std::make_shared(input->Dims(), output_dtype, input->GetDevice()); + auto device = input->GetDevice(); + const auto &stream = dynamic_cast( + infini_train::core::GetDeviceGuardImpl(device.type())->GetStream(device)) + ->maca_stream(); + const size_t num_elements = grad_output->NumElements(); + if (num_elements == 0) { + return grad_input; + } + const dim3 block(256); + const dim3 grid(std::min(CEIL_DIV(num_elements, block.x), static_cast(65535))); + + if (input->Dtype() == DataType::kFLOAT32 && grad_output->Dtype() == DataType::kFLOAT32) { + SwiGLUBackwardKernel<<>>( + static_cast(grad_input->DataPtr()), static_cast(input->DataPtr()), + static_cast(grad_output->DataPtr()), hidden, num_elements); + } else if (input->Dtype() == DataType::kBFLOAT16 && grad_output->Dtype() == DataType::kBFLOAT16) { + SwiGLUBackwardKernel<<>>(static_cast<__maca_bfloat16 *>(grad_input->DataPtr()), + static_cast(input->DataPtr()), + static_cast(grad_output->DataPtr()), + hidden, num_elements); + } else if (input->Dtype() == DataType::kBFLOAT16 && grad_output->Dtype() == DataType::kFLOAT32) { + SwiGLUBackwardKernel<<>>( + static_cast(grad_input->DataPtr()), static_cast(input->DataPtr()), + static_cast(grad_output->DataPtr()), hidden, num_elements); + } else if (input->Dtype() == DataType::kFLOAT32 && grad_output->Dtype() == DataType::kBFLOAT16) { + SwiGLUBackwardKernel<<>>( + static_cast(grad_input->DataPtr()), static_cast(input->DataPtr()), + static_cast(grad_output->DataPtr()), hidden, num_elements); + } else { + LOG_LOC(FATAL, "MACA SwiGLUBackward: unsupported data type combination"); + } + return grad_input; +} +} // namespace infini_train::kernels::maca + +#define REGISTER_MACA_SWIGLU_KERNEL(kernel_name) \ + REGISTER_KERNEL(infini_train::Device::DeviceType::kPrivateUse1, kernel_name, \ + infini_train::kernels::maca::kernel_name) + +REGISTER_MACA_SWIGLU_KERNEL(SwiGLUForward) +REGISTER_MACA_SWIGLU_KERNEL(SwiGLUBackward) + +#undef REGISTER_MACA_SWIGLU_KERNEL diff --git a/backends/maca/src/kernels/topk.maca b/backends/maca/src/kernels/topk.maca new file mode 100644 index 0000000..87f7148 --- /dev/null +++ b/backends/maca/src/kernels/topk.maca @@ -0,0 +1,157 @@ +#include +#include +#include + +#include "glog/logging.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 "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; +} + +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..9f3b281 --- /dev/null +++ b/backends/maca/src/kernels/transform.maca @@ -0,0 +1,599 @@ +#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 { + +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 { + +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..13bc259 --- /dev/null +++ b/backends/maca/src/kernels/vocab_parallel_cross_entropy.maca @@ -0,0 +1,130 @@ +#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 { + +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..31a1d7a --- /dev/null +++ b/backends/maca/src/runtime/maca_guard_impl.cc @@ -0,0 +1,533 @@ +#include "runtime/maca_guard_impl.h" + +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include + +#include "gflags/gflags.h" + +#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" + +DEFINE_bool(maca_multithread_workarounds, true, "Enable MACA multithread allocation, copy and example workarounds"); +DEFINE_bool(maca_retain_async_pool, false, "Retain pages in the MACA default async memory pool"); + +namespace infini_train::core::maca { +namespace { +constexpr size_t kBytesPerMB = 1024ULL * 1024ULL; + +// Read /proc/self/cmdline and return --tensor_parallel value, or 1 if absent +// or unparseable. +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; +} + +// Capture flags once so allocation and free always use the same implementation. +static bool g_multithread_workarounds = false; +static bool g_retain_async_pool = false; + +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; + +struct AllocationRecord { + int device_index = -1; + size_t size = 0; +}; + +struct AllocationStats { + size_t current = 0; + size_t high = 0; +}; + +// Synchronous mcMalloc allocations are not included in the default async +// mempool attributes, so track them separately and merge both peak values. +static std::mutex g_allocation_mutex; +static std::unordered_map g_allocations; +static std::vector g_allocation_stats; + +void TrackAllocation(void *ptr, int device_index, size_t size) { + if (ptr == nullptr) { + return; + } + + CHECK(device_index >= 0 && static_cast(device_index) < g_allocation_stats.size()) << std::format( + "MACA device index {} out of allocation stats range [0, {}).", device_index, g_allocation_stats.size()); + + std::lock_guard lock(g_allocation_mutex); + g_allocations[ptr] = AllocationRecord{device_index, size}; + auto &stats = g_allocation_stats.at(device_index); + stats.current += size; + stats.high = std::max(stats.high, stats.current); +} + +void UntrackAllocation(void *ptr) { + if (ptr == nullptr) { + return; + } + + std::lock_guard lock(g_allocation_mutex); + auto it = g_allocations.find(ptr); + if (it == g_allocations.end()) { + return; + } + + auto &stats = g_allocation_stats.at(it->second.device_index); + CHECK_GE(stats.current, it->second.size) << "MACA allocation tracker underflow."; + stats.current -= it->second.size; + g_allocations.erase(it); +} + +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(); + + if (g_retain_async_pool) { + // FIXME(cx): MACA 3.5.3.18 can fault when the default async pool releases + // pages to the driver and later allocations reuse that address range. + // Keep freed blocks resident; explicit mcMemPoolTrimTo can still release them. + mcMemPool_t pool; + MACA_CHECK(mcDeviceGetDefaultMemPool(&pool, device.index())); + uint64_t threshold = std::numeric_limits::max(); + MACA_CHECK(mcMemPoolSetAttribute(pool, mcMemPoolAttrReleaseThreshold, &threshold)); + } + + 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. + g_multithread_workarounds = FLAGS_maca_multithread_workarounds; + g_retain_async_pool = FLAGS_maca_retain_async_pool; + if (g_multithread_workarounds) { + setenv("MACA_LAUNCH_BLOCKING", "1", 0); + if (ReadTensorParallelFromCmdline() > 1) { + setenv("MCCL_P2P_DISABLE", "1", 0); + } + } + + const char *blocking = std::getenv("MACA_LAUNCH_BLOCKING"); + const char *p2p = std::getenv("MCCL_P2P_DISABLE"); + LOG(ERROR) << "MACA runtime: multithread_workarounds=" << (g_multithread_workarounds ? "true" : "false") + << " allocator=" << (g_multithread_workarounds ? "sync" : "async") + << " retain_async_pool=" << (g_retain_async_pool ? "true" : "false") + << " MACA_LAUNCH_BLOCKING=" << (blocking ? blocking : "unset") + << " MCCL_P2P_DISABLE=" << (p2p ? p2p : "unset"); + 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_allocation_stats.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) { + const mcError_t status = mcMalloc(dev_ptr, size); + if (status == mcErrorMemoryAllocation) { + int device_index = -1; + size_t free_bytes = 0, total_bytes = 0; + const mcError_t device_status = mcGetDevice(&device_index); + const mcError_t memory_status = mcMemGetInfo(&free_bytes, &total_bytes); + LOG(ERROR) << "MACA allocation failed: requested_bytes=" << size << " device=" << device_index + << " free_bytes=" << free_bytes << " total_bytes=" << total_bytes + << " device_query=" << mcGetErrorString(device_status) + << " memory_query=" << mcGetErrorString(memory_status); + } + MACA_CHECK(status); + + int device_index = -1; + MACA_CHECK(mcGetDevice(&device_index)); + TrackAllocation(*dev_ptr, device_index, size); +} + +void MacaGuardImpl::MallocAsync(void **dev_ptr, size_t size, Stream *stream) { + if (size == 0) { + *dev_ptr = nullptr; + return; + } + if (g_multithread_workarounds) { + Malloc(dev_ptr, size); + return; + } + auto maca_stream = GetMacaStream(stream); + const mcError_t status = mcMallocAsync(dev_ptr, size, maca_stream); + if (status == mcErrorMemoryAllocation) { + int device_index = -1; + size_t free_bytes = 0, total_bytes = 0; + const mcError_t device_status = mcGetDevice(&device_index); + const mcError_t memory_status = mcMemGetInfo(&free_bytes, &total_bytes); + LOG(ERROR) << "MACA async allocation failed: requested_bytes=" << size << " device=" << device_index + << " free_bytes=" << free_bytes << " total_bytes=" << total_bytes + << " retain_async_pool=" << (g_retain_async_pool ? "true" : "false") + << " device_query=" << mcGetErrorString(device_status) + << " memory_query=" << mcGetErrorString(memory_status); + } + MACA_CHECK(status); +} + +void MacaGuardImpl::Free(void *dev_ptr) { + UntrackAllocation(dev_ptr); + MACA_CHECK(mcFree(dev_ptr)); +} + +void MacaGuardImpl::FreeAsync(void *dev_ptr, Stream *stream) { + if (dev_ptr == nullptr) { + return; + } + auto maca_stream = GetMacaStream(stream); + if (g_multithread_workarounds) { + // A synchronous free must wait for work already submitted to this stream. + MACA_CHECK(mcStreamSynchronize(maca_stream)); + Free(dev_ptr); + return; + } + UntrackAllocation(dev_ptr); + MACA_CHECK(mcFreeAsync(dev_ptr, maca_stream)); +} + +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::unique_lock lock(g_memcpy_mutex, std::defer_lock); + if (g_multithread_workarounds) { + lock.lock(); + } + 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 { + CheckMacaDevice(device); + + int current_device = -1; + MACA_CHECK(mcGetDevice(¤t_device)); + + SetDevice(device); + mcMemPool_t pool; + MACA_CHECK(mcDeviceGetDefaultMemPool(&pool, device.index())); + + uint64_t zero = 0; + MACA_CHECK(mcMemPoolSetAttribute(pool, mcMemPoolAttrUsedMemHigh, &zero)); + MACA_CHECK(mcMemPoolSetAttribute(pool, mcMemPoolAttrReservedMemHigh, &zero)); + + { + std::lock_guard lock(g_allocation_mutex); + auto &stats = g_allocation_stats.at(device.index()); + stats.high = stats.current; + } + + MACA_CHECK(mcSetDevice(current_device)); +} + +std::pair MacaGuardImpl::GetMemPoolPeakMB(Device device) const { + CheckMacaDevice(device); + + int current_device = -1; + MACA_CHECK(mcGetDevice(¤t_device)); + + SetDevice(device); + mcMemPool_t pool; + MACA_CHECK(mcDeviceGetDefaultMemPool(&pool, device.index())); + + uint64_t used = 0; + MACA_CHECK(mcMemPoolGetAttribute(pool, mcMemPoolAttrUsedMemHigh, &used)); + + uint64_t reserved = 0; + MACA_CHECK(mcMemPoolGetAttribute(pool, mcMemPoolAttrReservedMemHigh, &reserved)); + + size_t tracked_high = 0; + { + std::lock_guard lock(g_allocation_mutex); + tracked_high = g_allocation_stats.at(device.index()).high; + } + + MACA_CHECK(mcSetDevice(current_device)); + + const size_t used_peak = tracked_high + static_cast(used); + const size_t reserved_peak = tracked_high + static_cast(reserved); + return {used_peak / kBytesPerMB, reserved_peak / kBytesPerMB}; +} + +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..ae3f96a --- /dev/null +++ b/backends/maca/src/runtime/maca_guard_impl.h @@ -0,0 +1,91 @@ +#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 { + +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/README.md b/scripts/README.md new file mode 100644 index 0000000..f5be6ef --- /dev/null +++ b/scripts/README.md @@ -0,0 +1,141 @@ +# 代码格式和静态检查 + +这些脚本及 `tools/clang_tidy` 插件迁移自 `third_party/InfiniTrain`， +以下命令供后续手动使用。Python 要求 3.10 或更新版本。 +相对的源码、构建目录、插件和报告路径均以本仓库根目录为基准。 + +## 迁移适配 + +| 项目 | Backends 中的行为 | +| --- | --- | +| 源码范围 | clang-tidy/IWYU 只选择 `backends/`、`tools/` 中属于本仓库的宿主编译单元；排除 `third_party`、构建目录和预期编译失败的 `*_compile_fail.cc` | +| 设备源码 | 格式化支持 `.cu`、`.maca`、`.mlu`、`.cl`；clang-tidy/IWYU 排除这些编译单元，并输出排除数量 | +| 头文件变更 | 检查编译数据库中全部宿主编译单元，覆盖间接引用；不再只选一个直接引用者 | +| 构建、依赖或检查规则变更 | 检查全部宿主编译单元；CMake 变更后需先重新生成编译数据库 | +| 编译数据库 | 派生过滤后的数据库，保留原编译参数；没有本仓库宿主条目或变更的宿主源文件缺少条目时均报错，需启用相关目标并重新生成数据库 | +| 格式化 | 沿用根目录 `.clang-format` 的 4 空格、120 列；修正失败返回码及 Black 成功输出被误判的问题，去除 `colorama` 依赖 | +| 自定义 tidy 规则 | 保留 InfiniTrain 的规则名和插件名，插件独立于后端构建，要求 LLVM/Clang 18 | + +## clang-format 与 Black + +`clang_format_diff.py` 要求 **clang-format 16.0.6**，默认查找 +`clang-format-16`，也可通过 `--clang-format /path/to/clang-format` 指定。 +默认只检查或格式化暂存区变更行，`--ref` 使用指定引用与 HEAD 的 merge-base +到 HEAD 的差异，`--all` 选择本仓库全部已跟踪的 C/C++ 和设备源码。 +依赖、构建目录及删除的文件均会跳过。 + +```bash +# 检查暂存区变更行；省略 --check 会写回格式化结果 +python3 scripts/clang_format_diff.py --check +python3 scripts/clang_format_diff.py --check --ref HEAD~1 +python3 scripts/clang_format_diff.py --check --all + +# 原有入口：检查整个文件，支持 C/C++、设备源码与 Python +python3 scripts/format.py --path backends tools scripts --check +``` + +`format.py` 保留 `--ref`、`--path`、`--check`、`--c`、`--py` 参数。 +不传 `--ref` 和 `--path` 时选择暂存区中的文件；`--ref` 与上述脚本一致； +`--path` 单独使用时递归选择文件，也包含未跟踪文件。 +`--c` 默认 `clang-format-16`，`--py` 默认 `black`，均支持二进制绝对路径。 +检查使用返回码判断结果，`--check` 不会写回文件。 + +暂存区模式与 `--ref` 模式都基于差异选择范围，工具读取工作区文件。 +使用时应让所选文件的工作区内容与被比较的版本一致，避免部分暂存导致行号偏移。 + +## 后端编译数据库 + +clang-tidy 和 IWYU 需要从仓库根目录配置 CMake，并在构建目录导出 `compile_commands.json`。 +仅在 InfiniTrain 子目录开启导出时，数据库可能只有子模块条目，无法检查后端。 +MACA 示例配置如下，其他构建依赖按 [MACA 指南](../backends/maca/README.zh-CN.md) 准备： + +```bash +cmake -S . -B build/lint -G Ninja \ + -DINFINITRAIN_BACKEND=maca \ + -DMACA_PATH=/opt/maca \ + -DCMAKE_EXPORT_COMPILE_COMMANDS=ON \ + -DINFINITRAIN_BACKENDS_BUILD_EXAMPLES=OFF \ + -DBUILD_TEST=OFF +``` + +这里使用后端的真实编译器及 SDK。不要沿用上游 CPU CI 的纯 clang 配置来 +代替后端配置。手动构建使用 `build`，自动模型测试使用 `build/maca`； +静态检查使用独立的 `build/lint`，模型运行脚本只清理自己的 `build/maca` 目录。 +多个后端需分别配置和检查；`--all` 只覆盖当前数据库中存在的宿主编译单元。 + +运行检查脚本时，会在 `BUILD_DIR/clang-tidy-host` 或 `BUILD_DIR/iwyu-host` +写入过滤后的数据库，原数据库保持完整。编译参数和包含路径不会被删改。 + +## clang-tidy + +`.clang-tidy` 启用两个规则：非模板、非 constexpr 成员函数应移出头文件 +（空函数体等情况豁免），以及枚举常量使用 `kFooBar` 命名。 +迁移保留了这些上游规则；首次检查可能发现已有代码违反规则。 + +自定义规则需要 `InfiniTrainTidy.so`，必须与 clang-tidy 来自同一套 LLVM/Clang 18 +开发环境。不要用 MACA 的 mxcc 构建此插件。典型的独立构建方式如下， +LLVM/Clang 的 CMake 路径需按实际安装位置调整： + +```bash +cmake -S tools/clang_tidy -B build/clang-tidy-plugin -G Ninja \ + -DCMAKE_BUILD_TYPE=Release \ + -DCMAKE_C_COMPILER=clang-18 \ + -DCMAKE_CXX_COMPILER=clang++-18 \ + -DLLVM_DIR=/usr/lib/llvm-18/lib/cmake/llvm \ + -DClang_DIR=/usr/lib/llvm-18/lib/cmake/clang +cmake --build build/clang-tidy-plugin + +python3 scripts/clang_tidy.py --all \ + --build-dir build/lint \ + --clang-tidy clang-tidy-18 \ + --run-clang-tidy run-clang-tidy-18 \ + --load-plugin build/clang-tidy-plugin/InfiniTrainTidy.so +``` + +不传 `--all` 时选择暂存区变更，或使用 `--ref HEAD~1`。 +仅宿主源文件变化时，直接对选中的编译单元运行 clang-tidy 并限制诊断到变更行； +纯删除没有新增行范围时检查该文件全部内容。头文件、CMake、子模块、规则或插件变化 +会改为全部宿主检查，因此也需要 `run-clang-tidy-18`。 +迁移后的增量模式不再依赖 `clang-tidy-diff`，避免将没有编译命令的头文件交给它。 + +`--clang-tidy` 指定的二进制也会传给全量 runner。 +`--jobs` 控制并发；全量扫描在 `BUILD_DIR/clang-tidy-report` 输出原始和去重诊断， +可通过 `--report-dir` 改位置。脚本不会应用 tidy 自动修复。 + +## IWYU + +需要支持 C++20 的 `include-what-you-use`、`iwyu_tool.py`，修复模式另需 +`fix_includes.py`；这些工具应来自同一套安装。上游使用的 LLVM 17 环境不能直接 +证明与当前 MACA SDK 兼容。`scripts/iwyu.imp` 保留标准库与 Eigen 的映射， +没有添加未经验证的 MACA SDK 映射。 + +```bash +python3 scripts/iwyu.py --all --build-dir build/lint +python3 scripts/iwyu.py --ref HEAD~1 --build-dir build/lint +# 只预览 include 修复 +python3 scripts/iwyu.py --all --build-dir build/lint --fix-dry-run +``` + +默认选择暂存区变化的宿主源文件；头文件、CMake、子模块或 IWYU 规则变化时选择 +全部宿主编译单元。默认只检查，发现 include 问题返回非零。 +`--fix` 才会改文件，修复范围限制在本仓库的 `backends`、`tools` 下， +默认保留头文件 include；`--fix-header-removals` 才允许删除它们。 +可用 `--iwyu-tool` 和 `--fix-tool` 指定工具路径。 + +两个静态检查入口都支持 `--has-relevant-changes [--ref REF]`： +有相关变更返回 0，没有返回 1，读取失败返回 2；此模式不运行分析器。 + +## 尚需实际验证的适配 + +MACA 将 `.maca` 作为 CXX 并传入 `-x maca`，其中包含设备关键字、内核启动语法和 +厂商内建函数。普通 clang-tidy/IWYU 无法按此语言解析，因此脚本不将其强制当普通 +C++ 检查。设备文件及仅由它们引用的头文件仍需 SDK 支持的分析方式。 + +即便是宿主 `.cc`，MACA SDK 头文件、编译器隐式包含路径及厂商参数与 LLVM 18/IWYU +前端的兼容性也需要实际检查。若报 unknown argument、缺少 SDK 头文件或内建类型， +应先核对 SDK 和分析前端，而不是忽略错误或删除语言参数来获得通过结果。 +编译数据库过滤不代表这些兼容性问题已解决。 + +未迁移上游的 clang-tidy/IWYU 自动 CI：它们只配置 CPU 构建，不满足此仓库的 +后端及 SDK 要求。现有格式化 CI 仍只检查 `backends`，显式安装 clang-format 16.0.6 +和 Black；新增分析器由开发者手动运行。 diff --git a/scripts/clang_format_diff.py b/scripts/clang_format_diff.py new file mode 100644 index 0000000..8b441c5 --- /dev/null +++ b/scripts/clang_format_diff.py @@ -0,0 +1,189 @@ +#!/usr/bin/env python3 +"""Check or format tracked C/C++ files or lines changed in a Git diff.""" + +from __future__ import annotations + +import argparse +import codecs +import re +import subprocess +import sys +from pathlib import Path + +from lint_common import ROOT, capture, diff_arguments, is_owned_path + +CLANG_FORMAT_VERSION = "16.0.6" +SOURCE_SUFFIXES = { + ".c", + ".cc", + ".cpp", + ".cxx", + ".cu", + ".cuh", + ".h", + ".hh", + ".hpp", + ".hxx", + ".maca", + ".mlu", + ".cl", +} +HUNK_RE = re.compile(r"^@@ .* \+(\d+)(?:,(\d+))? @@") + + +def run(command: list[str]) -> subprocess.CompletedProcess[str]: + return subprocess.run( + command, cwd=ROOT, capture_output=True, text=True, check=False + ) + + +def diff_text(ref: str | None) -> str: + return capture( + [ + "git", + "-c", + "core.quotePath=false", + "diff", + *diff_arguments(ref), + "--no-ext-diff", + "--no-textconv", + "--no-renames", + "--unified=0", + "--", + ".", + ":(exclude)**/third_party/**", + ":(exclude)third_party/**", + ":(exclude)build/**", + ] + ) + + +def changed_lines(diff: str) -> dict[Path, list[tuple[int, int]]]: + result: dict[Path, list[tuple[int, int]]] = {} + current: Path | None = None + for line in diff.splitlines(): + if line.startswith("+++ "): + name = line[4:] + if name.startswith('"'): + name = codecs.escape_decode(name[1:-1].encode("utf-8"))[0].decode( + "utf-8", errors="surrogateescape" + ) + current = None if name == "/dev/null" else Path(name.removeprefix("b/")) + continue + if current is None: + continue + match = HUNK_RE.match(line) + if match is None: + continue + start = int(match.group(1)) + count = int(match.group(2) or 1) + if count: + result.setdefault(current, []).append((start, start + count - 1)) + return result + + +def validate_version(binary: str) -> bool: + try: + process = run([binary, "--version"]) + except OSError as error: + print(f"error: cannot run formatter {binary}: {error}", file=sys.stderr) + return False + version = f"{process.stdout}\n{process.stderr}" + match = re.search(r"\bversion\s+(\d+\.\d+\.\d+)\b", version) + if ( + process.returncode != 0 + or match is None + or match.group(1) != CLANG_FORMAT_VERSION + ): + print( + f"error: clang-format {CLANG_FORMAT_VERSION} is required; " + f"got: {version.strip() or 'unknown'}", + file=sys.stderr, + ) + return False + return True + + +def validate_config(binary: str) -> bool: + process = run( + [ + binary, + "-style=file", + f"--assume-filename={ROOT / 'format_check.cc'}", + "-dump-config", + ] + ) + if process.returncode == 0: + return True + print("error: failed to load .clang-format", file=sys.stderr) + if process.stderr: + print(process.stderr, file=sys.stderr, end="") + return False + + +def format_file( + binary: str, path: Path, ranges: list[tuple[int, int]], check: bool +) -> bool: + command = [binary, "-style=file"] + command += ["--dry-run", "--Werror"] if check else ["-i"] + command += [f"--lines={start}:{end}" for start, end in ranges] + command.append(str(path)) + process = run(command) + if process.returncode == 0: + return True + print(f"error: clang-format failed for {path}", file=sys.stderr) + if process.stderr: + print(process.stderr, file=sys.stderr, end="") + return False + + +def main() -> int: + parser = argparse.ArgumentParser() + parser.add_argument("--check", action="store_true", help="do not modify files") + scope = parser.add_mutually_exclusive_group() + scope.add_argument("--ref", help="compare with the merge-base of REF") + scope.add_argument( + "--all", action="store_true", help="format all tracked project files" + ) + parser.add_argument("--clang-format", default="clang-format-16") + args = parser.parse_args() + + try: + if not validate_version(args.clang_format) or not validate_config( + args.clang_format + ): + return 2 + if args.all: + ranges = { + Path(name): [] + for name in capture(["git", "ls-files", "-z"]).split("\0") + if name and is_owned_path(Path(name)) + } + else: + ranges = changed_lines(diff_text(args.ref)) + except (OSError, RuntimeError, ValueError) as error: + print(f"error: {error}", file=sys.stderr) + return 2 + files = sorted( + path + for path in ranges + if is_owned_path(path) + and path.suffix.lower() in SOURCE_SUFFIXES + and (ROOT / path).is_file() + ) + if not files: + print("No supported C/C++ files require clang-format.") + return 0 + try: + results = [ + format_file(args.clang_format, path, ranges[path], args.check) + for path in files + ] + except OSError as error: + print(f"error: {error}", file=sys.stderr) + return 2 + return 0 if all(results) else 1 + + +if __name__ == "__main__": + raise SystemExit(main()) diff --git a/scripts/clang_tidy.py b/scripts/clang_tidy.py new file mode 100644 index 0000000..a54261f --- /dev/null +++ b/scripts/clang_tidy.py @@ -0,0 +1,327 @@ +#!/usr/bin/env python3 +"""Run the repository clang-tidy configuration on changed code or all TUs.""" + +from __future__ import annotations + +import argparse +import json +import os +import re +import shutil +import subprocess +import sys +from concurrent.futures import ThreadPoolExecutor, as_completed +from pathlib import Path + +from clang_format_diff import changed_lines, diff_text +from lint_common import ( + ROOT, + changed_paths, + header_filter_pattern, + host_entries, + needs_full_scan, + relevant_change, + root_path, + select_sources, + write_host_database, +) + +RELEVANT_FILES = { + ".gitmodules", + ".clang-tidy", + "CMakePresets.json", + ".github/workflows/clang-tidy.yaml", + "scripts/clang_tidy.py", + "scripts/clang_format_diff.py", + "scripts/lint_common.py", + "tools/clang_tidy/", +} +DIAGNOSTIC_RE = re.compile( + r"^(.*):(\d+):(\d+):\s+(warning|error|fatal error):\s+" + r"(.*?)(?:\s+\[([^\]]+)\])?\s*$" +) + + +def execute_captured(command: list[str]) -> tuple[int, str]: + print("+", " ".join(str(part) for part in command), flush=True) + process = subprocess.run( + command, + cwd=ROOT, + stdout=subprocess.PIPE, + stderr=subprocess.STDOUT, + text=True, + check=False, + ) + return process.returncode, process.stdout + + +def find_tool(explicit: str | None, candidates: list[str]) -> str | None: + if explicit: + return explicit + return next((path for name in candidates if (path := shutil.which(name))), None) + + +def validate_version(binary: str) -> None: + process = subprocess.run( + [binary, "--version"], cwd=ROOT, capture_output=True, text=True, check=False + ) + version = process.stdout + process.stderr + if ( + process.returncode != 0 + or re.search(r"\b(?:LLVM|clang) version 18\.", version, re.IGNORECASE) is None + ): + raise ValueError( + "clang-tidy 18 is required by the InfiniTrainTidy plugin; " + f"got: {version.strip() or 'unknown'}" + ) + + +def has_relevant_changes(ref: str | None) -> bool: + relevant = [ + path for path in changed_paths(ref) if relevant_change(path, RELEVANT_FILES) + ] + if not relevant: + print("No clang-tidy-relevant changes found.") + return False + print("clang-tidy-relevant changes:") + for path in relevant: + print(f" {path.as_posix()}") + return True + + +def plugin_arguments(plugin: Path) -> list[str]: + return ["-load", str(plugin)] + + +def display_path(path_text: str) -> str: + path = Path(path_text) + try: + if path.is_absolute(): + path = path.resolve().relative_to(ROOT) + except (OSError, ValueError): + pass + return path.as_posix() + + +def unique_diagnostics(output: str) -> tuple[list[str], int]: + unique: dict[tuple[str, str, str, str, str, str], int] = {} + for line in output.splitlines(): + match = DIAGNOSTIC_RE.match(line) + if match is None: + continue + path, line_number, column, severity, message, check_name = match.groups() + key = ( + display_path(path), + line_number, + column, + severity, + message, + check_name or "", + ) + unique[key] = unique.get(key, 0) + 1 + + lines = [] + ordered = sorted( + unique.items(), + key=lambda item: ( + item[0][0], + int(item[0][1]), + int(item[0][2]), + item[0][3:], + ), + ) + for key, count in ordered: + path, line_number, column, severity, message, check_name = key + suffix = f" [{check_name}]" if check_name else "" + repeated = f" (reported {count} times)" if count > 1 else "" + lines.append( + f"{path}:{line_number}:{column}: {severity}: {message}{suffix}{repeated}" + ) + return lines, sum(unique.values()) + + +def write_full_scan_reports(output: str, report_dir: Path) -> tuple[Path, Path]: + report_dir.mkdir(parents=True, exist_ok=True) + raw_path = report_dir / "clang-tidy-raw.txt" + unique_path = report_dir / "clang-tidy-unique.txt" + raw_path.write_text(output, encoding="utf-8") + + diagnostics, total = unique_diagnostics(output) + duplicate_count = total - len(diagnostics) + summary = ( + f"Unique diagnostics: {len(diagnostics)}; raw diagnostics: {total}; " + f"duplicates removed: {duplicate_count}." + ) + unique_path.write_text("\n".join([summary, *diagnostics, ""]), encoding="utf-8") + print(summary) + for diagnostic in diagnostics: + print(diagnostic) + print(f"Raw log: {raw_path}") + print(f"Deduplicated report: {unique_path}") + return raw_path, unique_path + + +def run_changed( + args: argparse.Namespace, clang_tidy: str, entries: list[dict[str, object]] +) -> int: + changes = changed_paths(args.ref) + if needs_full_scan(changes, RELEVANT_FILES): + print("Header, build, dependency or rule changes: scanning all host TUs.") + return run_all(args, clang_tidy, entries) + sources = select_sources(entries, changes, RELEVANT_FILES) + if not sources: + print("No host translation units require clang-tidy.") + return 0 + database = write_host_database(args.build_dir, "clang-tidy", entries) + ranges = changed_lines(diff_text(args.ref)) + + def check_source(source: Path) -> tuple[int, str]: + command = [ + clang_tidy, + "-p", + str(database), + "--config-file=" + str(ROOT / ".clang-tidy"), + "--header-filter=" + header_filter_pattern(), + "--quiet", + *plugin_arguments(args.load_plugin), + ] + lines = ranges.get(source.relative_to(ROOT)) + if lines: + command.append( + "--line-filter=" + json.dumps([{"name": str(source), "lines": lines}]) + ) + # Pure deletions have no new line range; check their complete host TU. + command.append(str(source)) + return execute_captured(command) + + failed = False + with ThreadPoolExecutor(max_workers=args.jobs) as executor: + futures = [executor.submit(check_source, source) for source in sources] + for future in as_completed(futures): + return_code, output = future.result() + print(output, end="" if output.endswith("\n") else "\n") + failed = failed or return_code != 0 + return 1 if failed else 0 + + +def run_all( + args: argparse.Namespace, clang_tidy: str, entries: list[dict[str, object]] +) -> int: + runner = find_tool( + args.run_clang_tidy, + ["run-clang-tidy-18", "run-clang-tidy.py", "run-clang-tidy"], + ) + if runner is None: + print("error: run-clang-tidy is required for full host scans", file=sys.stderr) + return 2 + command = [] + if Path(runner).suffix.lower() == ".py": + command = [sys.executable, "-W", "ignore::SyntaxWarning"] + command += [ + runner, + "-p", + str(write_host_database(args.build_dir, "clang-tidy", entries)), + "-clang-tidy-binary", + clang_tidy, + "-config-file", + str(ROOT / ".clang-tidy"), + "-header-filter", + header_filter_pattern(), + "-j", + str(args.jobs), + "-quiet", + ] + command += plugin_arguments(args.load_plugin) + return_code, output = execute_captured(command) + diagnostics, _ = unique_diagnostics(output) + write_full_scan_reports(output, args.report_dir) + if return_code != 0 and not diagnostics: + print(output, file=sys.stderr, end="" if output.endswith("\n") else "\n") + return return_code + + +def main() -> int: + parser = argparse.ArgumentParser() + parser.add_argument( + "--all", + action="store_true", + help="scan every project TU in the filtered compile database", + ) + parser.add_argument( + "--has-relevant-changes", + action="store_true", + help="report whether the diff contains changes that require clang-tidy", + ) + parser.add_argument( + "--ref", help="compare with the merge-base of REF (otherwise staged changes)" + ) + parser.add_argument("--build-dir", type=Path, default=Path("build/lint")) + parser.add_argument( + "--report-dir", + type=Path, + help="full-scan report directory (default: BUILD_DIR/clang-tidy-report)", + ) + parser.add_argument("--jobs", type=int, default=min(4, max(1, os.cpu_count() or 1))) + parser.add_argument("--clang-tidy") + parser.add_argument("--run-clang-tidy") + parser.add_argument( + "--load-plugin", + type=Path, + help="path to the InfiniTrain clang-tidy plugin", + ) + args = parser.parse_args() + + if args.all and args.ref: + parser.error("--ref cannot be used with --all; --all scans every project TU") + if args.jobs < 1: + parser.error("--jobs must be positive") + + try: + if args.has_relevant_changes: + return 0 if has_relevant_changes(args.ref) else 1 + except (OSError, RuntimeError, ValueError) as error: + print(f"error: {error}", file=sys.stderr) + return 2 + + if args.load_plugin is None: + print( + "error: --load-plugin is required when running clang-tidy", file=sys.stderr + ) + return 2 + + args.build_dir = root_path(args.build_dir) + args.load_plugin = root_path(args.load_plugin) + args.report_dir = ( + root_path(args.report_dir) + if args.report_dir is not None + else args.build_dir / "clang-tidy-report" + ) + + if not (args.build_dir / "compile_commands.json").exists(): + print( + f"error: {args.build_dir / 'compile_commands.json'} does not exist", + file=sys.stderr, + ) + return 2 + if not args.load_plugin.is_file(): + print(f"error: {args.load_plugin} does not exist", file=sys.stderr) + return 2 + clang_tidy = find_tool(args.clang_tidy, ["clang-tidy-18", "clang-tidy"]) + if clang_tidy is None: + print("error: clang-tidy is required", file=sys.stderr) + return 2 + try: + validate_version(clang_tidy) + entries = host_entries(args.build_dir) + return ( + run_all(args, clang_tidy, entries) + if args.all + else run_changed(args, clang_tidy, entries) + ) + except (OSError, RuntimeError, ValueError) as error: + print(f"error: {error}", file=sys.stderr) + return 2 + + +if __name__ == "__main__": + raise SystemExit(main()) diff --git a/scripts/format.py b/scripts/format.py new file mode 100644 index 0000000..da63726 --- /dev/null +++ b/scripts/format.py @@ -0,0 +1,142 @@ +#!/usr/bin/env python3 +"""Format complete files selected by path, staged changes or a Git reference.""" + +from __future__ import annotations + +import argparse +import os +import subprocess +import sys +from pathlib import Path + +from lint_common import ROOT, changed_paths, is_owned_path, root_path + +# Supported file types. +SUPPORTED_FILES = { + ".h": "c", + ".hh": "c", + ".hpp": "c", + ".hxx": "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: + kind = SUPPORTED_FILES.get(file.suffix.lower()) + if kind is None: + return True + binary = formatter[kind] + command = [binary] + if kind == "c": + command += ["-style=file"] + command += ["--dry-run", "--Werror"] if check else ["-i"] + elif check: + command += ["--check"] + command.append(str(file)) + try: + process = subprocess.run( + command, cwd=ROOT, capture_output=True, text=True, check=False + ) + except OSError as error: + print(f"error: cannot run {binary}: {error}", file=sys.stderr) + return False + if process.returncode != 0: + print(f"error: {binary} failed for {file}", file=sys.stderr) + print(process.stdout + process.stderr, file=sys.stderr, end="") + return False + if not check: + print(f"Formatted: {file.relative_to(ROOT)}") + return True + + +def git_added_files(): + """Get files selected by staged changes, using repository-relative Git paths.""" + for file in changed_paths(None): + yield root_path(file) + + +def git_modified_since_ref(ref): + """Get files changed from the merge-base of ref to HEAD.""" + for file in changed_paths(ref): + yield root_path(file) + + +def list_files(paths): + """Recursively get repository files under the specified paths.""" + for path in paths: + if not is_owned_path(path): + continue + if path.is_file(): + yield path + elif path.is_dir(): + for dirpath, subdirectories, filenames in os.walk(path): + subdirectories[:] = [ + name + for name in subdirectories + if is_owned_path(Path(dirpath) / name) + ] + for name in filenames: + file = Path(dirpath) / name + if is_owned_path(file): + yield file + else: + raise ValueError(f"not a file or directory: {path}") + + +def filter_in_path(file: Path, paths) -> bool: + """Determine whether the file is in one of the selected paths.""" + return any(file.is_relative_to(path) for path in paths) + + +def main() -> int: + parser = argparse.ArgumentParser(description=__doc__) + parser.add_argument("--ref", help="compare with the merge-base of REF") + parser.add_argument( + "--path", nargs="*", type=Path, help="repository-relative paths" + ) + parser.add_argument( + "--check", action="store_true", help="check without modifying files" + ) + parser.add_argument("--c", default="clang-format-16", help="C/C++ formatter binary") + parser.add_argument("--py", default="black", help="Python formatter binary") + args = parser.parse_args() + + try: + if args.path is not None: + args.path = [root_path(path) for path in args.path] + if args.ref is None and args.path is None: + files = git_added_files() + elif args.ref is None: + files = list_files(args.path) + elif args.path is None: + files = git_modified_since_ref(args.ref) + else: + 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 is_owned_path(file) or not file.is_file(): + continue + if not format_file(file, args.check, {"c": args.c, "py": args.py}): + formatted = False + return 0 if formatted else 1 + except (OSError, RuntimeError, ValueError) as error: + print(f"error: {error}", file=sys.stderr) + return 2 + + +if __name__ == "__main__": + raise SystemExit(main()) diff --git a/scripts/iwyu.imp b/scripts/iwyu.imp new file mode 100644 index 0000000..71d1056 --- /dev/null +++ b/scripts/iwyu.imp @@ -0,0 +1,7 @@ +[ + { "include": [ "@[\"<]bits/chrono\\.h[\">]", "private", "", "public" ] }, + { "include": [ "@[\"<]bits/std_abs\\.h[\">]", "private", "", "public" ] }, + { "symbol": [ "std::milli", "private", "", "public" ] }, + { "include": [ "@[\"<]Eigen/src/(Core|LU|Cholesky|SVD|QR|Eigenvalues|misc)/.*", "private", "\"Eigen/Dense\"", "public" ] }, + { "include": [ "@[\"<]Eigen/(Core|LU|Cholesky|SVD|QR|Eigenvalues)[\">]", "private", "\"Eigen/Dense\"", "public" ] } +] diff --git a/scripts/iwyu.py b/scripts/iwyu.py new file mode 100644 index 0000000..b48ecd7 --- /dev/null +++ b/scripts/iwyu.py @@ -0,0 +1,231 @@ +#!/usr/bin/env python3 +"""Run IWYU on changed translation units or the complete project database.""" + +from __future__ import annotations + +import argparse +import os +import re +import shutil +import subprocess +import sys +from pathlib import Path + +from lint_common import ( + ROOT, + changed_paths, + host_entries, + project_file_pattern, + relevant_change, + root_path, + select_sources, + write_host_database, +) + +RELEVANT_FILES = { + ".github/workflows/iwyu.yaml", + ".gitmodules", + "CMakePresets.json", + "scripts/iwyu.imp", + "scripts/iwyu.py", + "scripts/lint_common.py", +} +CORRECT_RESULT_RE = re.compile(r"^\(.+ has correct #includes/fwd-decls\)$") + + +def run_check(command: list[str]) -> int: + """Run IWYU while suppressing non-actionable success messages.""" + process = subprocess.Popen( + command, + cwd=ROOT, + stdout=subprocess.PIPE, + stderr=subprocess.STDOUT, + text=True, + bufsize=1, + ) + assert process.stdout is not None + pending_blank = False + emitted_output = False + for line in process.stdout: + if CORRECT_RESULT_RE.fullmatch(line.strip()): + continue + if not line.strip(): + pending_blank = True + continue + if pending_blank and emitted_output: + print() + print(line, end="", flush=True) + pending_blank = False + emitted_output = True + return process.wait() + + +def main() -> int: + parser = argparse.ArgumentParser() + mode = parser.add_mutually_exclusive_group() + mode.add_argument("--all", action="store_true", help="scan every project TU") + mode.add_argument( + "--has-relevant-changes", + action="store_true", + help="return 0 when the selected diff should trigger a full IWYU scan", + ) + parser.add_argument("--ref", help="compare with the merge-base of REF") + parser.add_argument("--build-dir", type=Path, default=Path("build/lint")) + parser.add_argument("--jobs", type=int, default=min(4, max(1, os.cpu_count() or 1))) + parser.add_argument("--iwyu-tool") + fix_mode = parser.add_mutually_exclusive_group() + fix_mode.add_argument( + "--fix", action="store_true", help="apply IWYU output with fix_includes.py" + ) + fix_mode.add_argument( + "--fix-dry-run", + action="store_true", + help="print the fix_includes.py diff without modifying files", + ) + parser.add_argument( + "--fix-header-removals", + action="store_true", + help="allow fix_includes.py to remove includes from header files", + ) + parser.add_argument("--fix-tool") + args = parser.parse_args() + args.build_dir = root_path(args.build_dir) + + if args.all and args.ref: + parser.error("--ref cannot be used with --all; --all scans every project TU") + if args.jobs < 1: + parser.error("--jobs must be positive") + if args.has_relevant_changes and (args.fix or args.fix_dry_run): + parser.error("--has-relevant-changes cannot be combined with a fix mode") + if args.fix_header_removals and not (args.fix or args.fix_dry_run): + parser.error("--fix-header-removals requires --fix or --fix-dry-run") + if args.fix_tool and not (args.fix or args.fix_dry_run): + parser.error("--fix-tool requires --fix or --fix-dry-run") + + if args.has_relevant_changes: + try: + relevant = [ + path + for path in changed_paths(args.ref) + if relevant_change(path, RELEVANT_FILES) + ] + except (OSError, RuntimeError, ValueError) as error: + print(f"error: {error}", file=sys.stderr) + return 2 + if relevant: + print("Relevant IWYU changes:") + for path in relevant: + print(f" {path}") + return 0 + print("No relevant changes require IWYU.") + return 1 + + tool = args.iwyu_tool or shutil.which("iwyu_tool.py") or shutil.which("iwyu_tool") + if tool is None: + print("error: iwyu_tool is required", file=sys.stderr) + return 2 + fixer = None + if args.fix or args.fix_dry_run: + fixer = ( + args.fix_tool + or shutil.which("fix_includes.py") + or shutil.which("fix_include") + ) + if fixer is None: + print( + "error: fix_includes.py or fix_include is required for fix mode", + file=sys.stderr, + ) + return 2 + try: + entries = host_entries(args.build_dir) + sources = ( + sorted({Path(str(entry["file"])) for entry in entries}) + if args.all + else select_sources(entries, changed_paths(args.ref), RELEVANT_FILES) + ) + database = ( + write_host_database(args.build_dir, "iwyu", entries) if sources else None + ) + except (OSError, RuntimeError, ValueError) as error: + print(f"error: {error}", file=sys.stderr) + return 2 + if not sources: + print("No host translation units require IWYU.") + return 0 + + command = [ + tool, + "-j", + str(args.jobs), + "-p", + str(database), + *[str(source) for source in sources], + "--", + "-Xiwyu", + "--cxx17ns", + "-Xiwyu", + "--mapping_file=" + str(ROOT / "scripts/iwyu.imp"), + ] + if not (args.fix or args.fix_dry_run): + command.extend(["-Xiwyu", "--error=1"]) + display_command = [ + tool, + "-j", + str(args.jobs), + "-p", + str(database), + f"<{len(sources)} translation units>", + "--", + *command[command.index("--") + 1 :], + ] + print("+", " ".join(display_command), flush=True) + if not (args.fix or args.fix_dry_run): + return run_check(command) + + analysis = subprocess.run( + command, + cwd=ROOT, + stdout=subprocess.PIPE, + stderr=subprocess.STDOUT, + text=True, + check=False, + ) + print(analysis.stdout, end="") + if analysis.returncode != 0: + print( + "error: IWYU analysis failed; no automatic fixes were applied", + file=sys.stderr, + ) + return analysis.returncode + + fix_command = [ + fixer, + "--basedir", + str(ROOT), + "--only_re", + project_file_pattern(), + "--nocomments", + "--blank_lines", + "--noreorder", + "--nosafe_headers" if args.fix_header_removals else "--safe_headers", + ] + if args.fix_dry_run: + fix_command.append("--dry_run") + print("+", " ".join(fix_command), flush=True) + fixed = subprocess.run( + fix_command, + cwd=ROOT, + input=analysis.stdout, + text=True, + check=False, + ) + if args.fix_dry_run: + print("IWYU fix preview completed; no files were modified.") + elif fixed.returncode == 0: + print("IWYU fixes applied. Run formatting, build, and tests before committing.") + return fixed.returncode + + +if __name__ == "__main__": + raise SystemExit(main()) diff --git a/scripts/lint_common.py b/scripts/lint_common.py new file mode 100644 index 0000000..516d0e3 --- /dev/null +++ b/scripts/lint_common.py @@ -0,0 +1,182 @@ +"""Repository scope and host compile-database helpers for the lint drivers.""" + +from __future__ import annotations + +import json +import re +import subprocess +from pathlib import Path + +ROOT = Path(__file__).resolve().parent.parent +PROJECT_ROOTS = {"backends", "tools"} +HOST_SOURCES = {".c", ".cc", ".cpp", ".cxx"} +DEVICE_SOURCES = {".cu", ".maca", ".mlu", ".cl"} +HEADERS = {".h", ".hh", ".hpp", ".hxx", ".cuh"} +EXCLUDED_DIRS = {".git", "third_party", "build", ".cache", "__pycache__", ".venv"} +BUILD_PREFIXES = ("build-", "build_", "cmake-build-") + + +def root_path(path: Path) -> Path: + return (ROOT / path).resolve() + + +def is_owned_path(path: Path) -> bool: + try: + relative = root_path(path).relative_to(ROOT) + except ValueError: + return False + return not any( + part in EXCLUDED_DIRS or part.startswith(BUILD_PREFIXES) + for part in relative.parts + ) + + +def is_project_path(path: Path) -> bool: + if not is_owned_path(path): + return False + relative = root_path(path).relative_to(ROOT) + return bool(relative.parts) and relative.parts[0] in PROJECT_ROOTS + + +def capture(command: list[str]) -> str: + process = subprocess.run( + command, cwd=ROOT, capture_output=True, text=True, check=False + ) + if process.returncode != 0: + raise RuntimeError(process.stderr.strip() or "command failed") + return process.stdout + + +def diff_arguments(ref: str | None) -> list[str]: + if ref: + base = capture(["git", "merge-base", ref, "HEAD"]).strip() + return [f"{base}..HEAD"] + return ["--cached"] + + +def changed_paths(ref: str | None) -> list[Path]: + output = capture( + [ + "git", + "diff", + *diff_arguments(ref), + "--no-renames", + "--name-only", + "--diff-filter=ACDM", + "-z", + "--", + ] + ) + return [Path(name) for name in output.split("\0") if name] + + +def needs_full_scan(changes: list[Path], relevant_files: set[str]) -> bool: + for path in changes: + normalized = path.as_posix() + if normalized in relevant_files or any( + normalized.startswith(prefix) + for prefix in relevant_files + if prefix.endswith("/") + ): + return True + if path.parts and path.parts[0] == "third_party": + return True + if not is_owned_path(path): + continue + if path.name == "CMakeLists.txt" or path.suffix.lower() == ".cmake": + return True + if is_project_path(path) and path.suffix.lower() in HEADERS: + return True + return False + + +def relevant_change(path: Path, relevant_files: set[str]) -> bool: + return needs_full_scan([path], relevant_files) or ( + is_project_path(path) and path.suffix.lower() in HOST_SOURCES | DEVICE_SOURCES + ) + + +def host_entries(build_dir: Path) -> list[dict[str, object]]: + database = json.loads( + (build_dir / "compile_commands.json").read_text(encoding="utf-8") + ) + if not isinstance(database, list): + raise ValueError("compile_commands.json must contain a list") + filtered = [] + devices = set() + for entry in database: + if not isinstance(entry, dict) or "file" not in entry: + raise ValueError("invalid compile database entry: missing file") + directory = (build_dir / str(entry.get("directory", build_dir))).resolve() + source = (directory / str(entry["file"])).resolve() + if not is_project_path(source): + continue + if source.suffix.lower() in DEVICE_SOURCES: + devices.add(source) + if source.suffix.lower() not in HOST_SOURCES or source.name.endswith( + "_compile_fail.cc" + ): + continue + filtered.append({**entry, "directory": str(directory), "file": str(source)}) + print( + f"Host compile commands: {len(filtered)}; " + f"device translation units excluded: {len(devices)} (.cu/.maca/.mlu/.cl).", + flush=True, + ) + if not filtered: + raise ValueError( + "No backends/tools host translation units in compile_commands.json. " + "Configure the repository root with -DINFINITRAIN_BACKEND= " + "and -DCMAKE_EXPORT_COMPILE_COMMANDS=ON (see scripts/README.md)." + ) + return filtered + + +def write_host_database( + build_dir: Path, tool: str, entries: list[dict[str, object]] +) -> Path: + output_dir = build_dir / f"{tool}-host" + output_dir.mkdir(parents=True, exist_ok=True) + (output_dir / "compile_commands.json").write_text( + json.dumps(entries, indent=2) + "\n", encoding="utf-8" + ) + return output_dir + + +def select_sources( + entries: list[dict[str, object]], changes: list[Path], relevant_files: set[str] +) -> list[Path]: + available = {Path(str(entry["file"])) for entry in entries} + changed_sources = { + root_path(path) + for path in changes + if is_project_path(path) + and path.suffix.lower() in HOST_SOURCES + and not path.name.endswith("_compile_fail.cc") + and root_path(path).is_file() + } + missing = changed_sources - available + if missing: + paths = ", ".join(str(path.relative_to(ROOT)) for path in sorted(missing)) + raise ValueError( + f"Changed host sources missing from compile_commands.json: {paths}. " + "Regenerate the database with the relevant backend/targets enabled." + ) + if needs_full_scan(changes, relevant_files): + return sorted(available) + return sorted(available.intersection(root_path(path) for path in changes)) + + +def project_file_pattern() -> str: + roots = "|".join(re.escape(root) for root in sorted(PROJECT_ROOTS)) + excluded = "|".join(re.escape(name) for name in sorted(EXCLUDED_DIRS)) + excluded += "|" + "|".join(re.escape(prefix) + "[^/]*" for prefix in BUILD_PREFIXES) + repository = re.escape(ROOT.as_posix()) + return rf"^(?:{repository}/|\./)?(?!(?:[^/]+/)*(?:{excluded})(?:/|$))(?:{roots})/" + + +def header_filter_pattern() -> str: + # clang-tidy uses LLVM's POSIX regex engine, not Python's re engine. + roots = "|".join(re.escape(root) for root in sorted(PROJECT_ROOTS)) + repository = re.escape(ROOT.as_posix()) + return rf"^({repository}/|[.]/)?({roots})/" diff --git a/third_party/InfiniTrain b/third_party/InfiniTrain new file mode 160000 index 0000000..c90338c --- /dev/null +++ b/third_party/InfiniTrain @@ -0,0 +1 @@ +Subproject commit c90338ceb2dc528816bf78c422f1f5483f2be1ae diff --git a/tools/clang_tidy/CMakeLists.txt b/tools/clang_tidy/CMakeLists.txt new file mode 100644 index 0000000..3508b07 --- /dev/null +++ b/tools/clang_tidy/CMakeLists.txt @@ -0,0 +1,35 @@ +cmake_minimum_required(VERSION 3.20) + +project(InfiniTrainTidy LANGUAGES C CXX) + +find_package(LLVM REQUIRED CONFIG) +find_package(Clang REQUIRED CONFIG) + +if(NOT LLVM_PACKAGE_VERSION VERSION_GREATER_EQUAL 18 OR + LLVM_PACKAGE_VERSION VERSION_GREATER_EQUAL 19) + message(FATAL_ERROR "InfiniTrainTidy requires LLVM/Clang 18") +endif() + +separate_arguments(LLVM_DEFINITIONS_LIST NATIVE_COMMAND "${LLVM_DEFINITIONS}") + +add_library(InfiniTrainTidy MODULE + InfiniTrainTidyModule.cpp + NoMemberDefinitionsInHeadersCheck.cpp +) + +target_compile_features(InfiniTrainTidy PRIVATE cxx_std_17) +target_include_directories(InfiniTrainTidy SYSTEM PRIVATE + ${LLVM_INCLUDE_DIRS} + ${CLANG_INCLUDE_DIRS} +) +target_compile_options(InfiniTrainTidy PRIVATE ${LLVM_DEFINITIONS_LIST}) +target_link_libraries(InfiniTrainTidy PRIVATE + clangAST + clangASTMatchers + clangBasic + clangTidy +) + +set_target_properties(InfiniTrainTidy PROPERTIES + PREFIX "" +) diff --git a/tools/clang_tidy/InfiniTrainTidyModule.cpp b/tools/clang_tidy/InfiniTrainTidyModule.cpp new file mode 100644 index 0000000..39fe5d6 --- /dev/null +++ b/tools/clang_tidy/InfiniTrainTidyModule.cpp @@ -0,0 +1,22 @@ +#include "NoMemberDefinitionsInHeadersCheck.h" +#include "clang-tidy/ClangTidyModule.h" +#include "clang-tidy/ClangTidyModuleRegistry.h" + +namespace clang::tidy::infinitrain { +namespace { + +class InfiniTrainTidyModule : public ClangTidyModule { +public: + void addCheckFactories(ClangTidyCheckFactories &factories) override { + factories.registerCheck("infinitrain-no-member-definitions-in-headers"); + } +}; + +static ClangTidyModuleRegistry::Add module("infinitrain-module", + "Adds InfiniTrain project checks."); + +} // namespace + +volatile int InfiniTrainTidyModuleAnchorSource = 0; + +} // namespace clang::tidy::infinitrain diff --git a/tools/clang_tidy/NoMemberDefinitionsInHeadersCheck.cpp b/tools/clang_tidy/NoMemberDefinitionsInHeadersCheck.cpp new file mode 100644 index 0000000..d394a46 --- /dev/null +++ b/tools/clang_tidy/NoMemberDefinitionsInHeadersCheck.cpp @@ -0,0 +1,78 @@ +#include "NoMemberDefinitionsInHeadersCheck.h" + +#include "clang-tidy/utils/FileExtensionsUtils.h" +#include "clang/AST/DeclCXX.h" +#include "clang/AST/Stmt.h" +#include "clang/ASTMatchers/ASTMatchFinder.h" +#include "llvm/ADT/StringRef.h" +#include "llvm/Support/Casting.h" + +namespace clang::tidy::infinitrain { +namespace { + +bool isTemplated(const CXXMethodDecl &method) { + if (method.getTemplatedKind() != FunctionDecl::TK_NonTemplate) { + return true; + } + + for (const DeclContext *context = method.getDeclContext(); context != nullptr; context = context->getParent()) { + const auto *record = dyn_cast(context); + if (record != nullptr + && (record->getDescribedClassTemplate() != nullptr || isa(record))) { + return true; + } + } + return false; +} + +} // namespace + +NoMemberDefinitionsInHeadersCheck::NoMemberDefinitionsInHeadersCheck(StringRef name, ClangTidyContext *context) + : ClangTidyCheck(name, context), ignore_empty_bodies_(Options.get("IgnoreEmptyBodies", true)), + header_file_extensions_(context->getHeaderFileExtensions()) {} + +bool NoMemberDefinitionsInHeadersCheck::isLanguageVersionSupported(const LangOptions &lang_opts) const { + return lang_opts.CPlusPlus; +} + +void NoMemberDefinitionsInHeadersCheck::storeOptions(ClangTidyOptions::OptionMap &options) { + Options.store(options, "IgnoreEmptyBodies", ignore_empty_bodies_); +} + +void NoMemberDefinitionsInHeadersCheck::registerMatchers(ast_matchers::MatchFinder *finder) { + using namespace ast_matchers; + finder->addMatcher(cxxMethodDecl(isDefinition(), unless(isImplicit())).bind("method"), this); +} + +void NoMemberDefinitionsInHeadersCheck::check(const ast_matchers::MatchFinder::MatchResult &result) { + const auto *method = result.Nodes.getNodeAs("method"); + if (method == nullptr || method->isDefaulted() || method->isDeleted() || method->isConstexpr() + || method->getParent()->isLambda() || isTemplated(*method)) { + return; + } + + const auto *body = dyn_cast_or_null(method->getBody()); + if (ignore_empty_bodies_ && body != nullptr && body->body_empty()) { + return; + } + + SourceLocation location = method->getLocation(); + if (location.isInvalid() || location.isMacroID()) { + return; + } + if (const Stmt *body = method->getBody(); body != nullptr && body->getBeginLoc().isMacroID()) { + return; + } + + SourceManager &source_manager = *result.SourceManager; + location = source_manager.getSpellingLoc(location); + if (source_manager.isInSystemHeader(location) + || !utils::isSpellingLocInHeaderFile(location, source_manager, header_file_extensions_)) { + return; + } + + diag(location, "member function definition must be moved from the header to an " + "implementation file"); +} + +} // namespace clang::tidy::infinitrain diff --git a/tools/clang_tidy/NoMemberDefinitionsInHeadersCheck.h b/tools/clang_tidy/NoMemberDefinitionsInHeadersCheck.h new file mode 100644 index 0000000..fc2b61f --- /dev/null +++ b/tools/clang_tidy/NoMemberDefinitionsInHeadersCheck.h @@ -0,0 +1,22 @@ +#pragma once + +#include "clang-tidy/ClangTidyCheck.h" +#include "clang-tidy/FileExtensionsSet.h" + +namespace clang::tidy::infinitrain { + +class NoMemberDefinitionsInHeadersCheck : public ClangTidyCheck { +public: + NoMemberDefinitionsInHeadersCheck(StringRef name, ClangTidyContext *context); + + void registerMatchers(ast_matchers::MatchFinder *finder) override; + void check(const ast_matchers::MatchFinder::MatchResult &result) override; + bool isLanguageVersionSupported(const LangOptions &lang_opts) const override; + void storeOptions(ClangTidyOptions::OptionMap &options) override; + +private: + const bool ignore_empty_bodies_; + const FileExtensionsSet header_file_extensions_; +}; + +} // namespace clang::tidy::infinitrain