diff --git a/.gitattributes b/.gitattributes new file mode 100644 index 0000000..7753682 --- /dev/null +++ b/.gitattributes @@ -0,0 +1,5 @@ +* text=auto eol=lf +*.sh text eol=lf +*.mlu text eol=lf +*.cc text eol=lf +*.h text eol=lf diff --git a/.gitignore b/.gitignore index eeba729..02054ca 100644 --- a/.gitignore +++ b/.gitignore @@ -7,3 +7,11 @@ *.log *.report.rank* *.records.log.rank* + +# MLU provider local artifacts +*.cndump +core +core.* +__pycache__/ +/backends/mlu/tests/math_selftest/math_selftest +/backends/mlu/tests/math_selftest/*.o diff --git a/README.md b/README.md index 9220950..4b9541d 100644 --- a/README.md +++ b/README.md @@ -27,6 +27,7 @@ process may register at most one provider for `DeviceType::kPrivateUse1`. | Backend | Runtime | Collectives | Model examples | | ------- | ------- | ----------- | -------------- | | [MACA](backends/maca/README.md) | MACA | MCCL (optional) | GPT-2, LLaMA 3, Mixtral | +| [MLU](backends/mlu/README.md) (work in progress) | CNRT | — | GPT-2 | The model sources normally come from the pinned InfiniTrain submodule. Each backend supplies the runtime, kernels, collective implementation, and build @@ -114,6 +115,23 @@ CMake cannot replace a compiler after a language has been enabled. See [formatting and static checks](scripts/README.md) for clang-format, clang-tidy and IWYU setup, scope and manual usage (in Chinese). +## Cambricon MLU provider (work in progress) + +`backends/mlu` registers a `PrivateUse1` backend named `mlu` for Cambricon MLU +devices. Small-shape FP32 operators and the device math functions are covered by +regression tools; device GEMM, full GPT-2 10-step accuracy, BF16 and multi-card +collectives are **not** done yet. + +```bash +cmake -S . -B build-mlu -G Ninja \ + -DINFINITRAIN_BACKEND=mlu -DUSE_CUDA=OFF -DBUILD_TEST=OFF +cmake --build build-mlu --target mlu_gpt2 mlu_op_test -j +``` + +See [the MLU provider guide](backends/mlu/README.md) for the SDK requirements and +[its validation status](backends/mlu/docs/validation.md) for what has actually +been run. + ## License InfiniTrain Backends is released under the [MIT License](LICENSE). diff --git a/README.zh-CN.md b/README.zh-CN.md index 27b855e..9d090ea 100644 --- a/README.zh-CN.md +++ b/README.zh-CN.md @@ -25,6 +25,7 @@ InfiniTrain 子模块。每个构建目录只选择一个后端，每个进程 | 后端 | 运行时 | 集合通信 | 模型示例 | | ---- | ------ | -------- | -------- | | [MACA](backends/maca/README.zh-CN.md) | MACA | MCCL（可选） | GPT-2、LLaMA 3、Mixtral | +| [MLU](backends/mlu/README.md)（开发中） | CNRT | — | GPT-2 | 模型源码通常来自锁定版本的 InfiniTrain 子模块。各后端提供运行时、算子、集合通信 实现和构建集成，以便在对应加速器上运行其支持的模型示例。环境配置和使用方式见上表 @@ -104,6 +105,21 @@ CMake 无法在启用某种语言后更换其编译器。 clang-format、clang-tidy 和 IWYU 的范围、依赖及手动使用方式见 [代码格式和静态检查](scripts/README.md)。 +## 寒武纪 MLU 后端（开发中） + +`backends/mlu` 通过 `PrivateUse1` 注册名为 `mlu` 的后端。小 shape 的 FP32 算子 +和设备数学函数已有回归工具覆盖；设备 GEMM、完整 GPT-2 十步精度、BF16 和多卡 +集合通信**尚未完成**。 + +```bash +cmake -S . -B build-mlu -G Ninja \ + -DINFINITRAIN_BACKEND=mlu -DUSE_CUDA=OFF -DBUILD_TEST=OFF +cmake --build build-mlu --target mlu_gpt2 mlu_op_test -j +``` + +SDK 依赖见 [MLU 后端指南](backends/mlu/README.md)，实际验证状态见 +[验证状态](backends/mlu/docs/validation.md)。 + ## 许可证 InfiniTrain Backends 使用 [MIT 许可证](LICENSE)。 diff --git a/backends/mlu/CMakeLists.txt b/backends/mlu/CMakeLists.txt new file mode 100644 index 0000000..0f8fc26 --- /dev/null +++ b/backends/mlu/CMakeLists.txt @@ -0,0 +1,136 @@ +# Cambricon MLU provider. +# +# Registers a PrivateUse1 backend named "mlu": the runtime, the kernel host +# integration and the BANG C device code all live under this directory, so no +# InfiniTrain source file is modified. BANG C sources are compiled by cncc with +# a custom command; host sources are ordinary C++ compiled by the selected host +# compiler. + +find_program(CNCC_EXECUTABLE cncc HINTS "${NEUWARE_HOME}/bin" REQUIRED) +find_library(CNRT_LIBRARY cnrt HINTS "${NEUWARE_HOME}/lib64" REQUIRED) + +# CNNL ships separately from the base Neuware image. When it is missing the +# build still produces a working runtime and kernels; only the BLAS handle is +# left out. No operator calls CNNL yet - GEMM still uses the host fallback. +option(INFINITRAIN_MLU_WITH_CNNL "Compile the CNNL BLAS handle support" ON) +if(INFINITRAIN_MLU_WITH_CNNL) + find_library(CNNL_LIBRARY cnnl HINTS "${NEUWARE_HOME}/lib64") + if(NOT CNNL_LIBRARY OR NOT EXISTS "${NEUWARE_HOME}/include/cnnl.h") + message(WARNING + "CNNL was not found under ${NEUWARE_HOME}; configuring without BLAS " + "handle support. Install CNNL and reconfigure to enable it.") + set(INFINITRAIN_MLU_WITH_CNNL OFF) + endif() +endif() + +# BANG C target architecture; empty lets cncc use its default. +set(MLU_ARCH "" CACHE STRING "BANG C target arch, for example mtp_592") + +option(INFINITRAIN_MLU_BUILD_TEST_TOOLS + "Build the MLU operator and device-math regression tools" ON) + +# ------------------------------------------------------------------------------ +# Runtime, registration entry and kernel host integration (host compiler) +# ------------------------------------------------------------------------------ +set(MLU_BACKEND_SOURCES + src/backend.cc + src/runtime/mlu_runtime_common.cc + src/runtime/mlu_guard_impl.cc + src/kernels/no_op.cc + src/kernels/fill.cc + src/kernels/cast.cc + src/kernels/accumulate_grad.cc + src/kernels/elementwise.cc + src/kernels/embedding.cc + src/kernels/layernorm.cc + src/kernels/softmax.cc + src/kernels/cross_entropy.cc + src/kernels/transform.cc + src/kernels/linear.cc +) + +add_library(infini_train_mlu STATIC ${MLU_BACKEND_SOURCES}) +add_library(InfiniTrain::Backend::MLU ALIAS infini_train_mlu) + +# Only include/infini_train_mlu is public: examples and tools consume the +# registration entry point, never the provider's internal layout. Internal +# headers (common/, runtime/, kernels/) resolve against src/, which stays +# PRIVATE and is added explicitly by the targets that need it. +target_include_directories(infini_train_mlu + PUBLIC + "$" + PRIVATE + "${CMAKE_CURRENT_SOURCE_DIR}/src" + "${INFINITRAIN_SOURCE_DIR}" + "${NEUWARE_HOME}/include" + "${CMAKE_CURRENT_SOURCE_DIR}/third_party/bang_host" +) + +target_link_libraries(infini_train_mlu PUBLIC + InfiniTrain::infini_train + ${CNRT_LIBRARY} +) + +if(INFINITRAIN_MLU_WITH_CNNL) + target_compile_definitions(infini_train_mlu PUBLIC USE_CNNL=1) + target_link_libraries(infini_train_mlu PUBLIC ${CNNL_LIBRARY}) +endif() + +# ------------------------------------------------------------------------------ +# BANG C kernels (*.mlu): CMake has no BANG C language, so cncc is invoked +# directly and the resulting objects form their own archive. +# ------------------------------------------------------------------------------ +file(GLOB MLU_KERNEL_SOURCES CONFIGURE_DEPENDS + "${CMAKE_CURRENT_SOURCE_DIR}/src/kernels/*.mlu") + +set(MLU_KERNEL_OBJECTS) +foreach(src ${MLU_KERNEL_SOURCES}) + get_filename_component(name ${src} NAME_WE) + set(obj "${CMAKE_CURRENT_BINARY_DIR}/${name}.o") + set(arch_flags) + if(MLU_ARCH) + list(APPEND arch_flags --bang-mlu-arch=${MLU_ARCH}) + endif() + add_custom_command( + OUTPUT ${obj} + COMMAND ${CMAKE_COMMAND} -E env NEUWARE_HOME=${NEUWARE_HOME} + ${CNCC_EXECUTABLE} -std=c++17 -fPIC -c ${src} -o ${obj} ${arch_flags} + "-I${CMAKE_CURRENT_SOURCE_DIR}/src" + "-I${NEUWARE_HOME}/include" + DEPENDS ${src} + "${CMAKE_CURRENT_SOURCE_DIR}/src/kernels/mlu_kernel_api.h" + COMMENT "cncc ${name}.mlu" + VERBATIM + ) + list(APPEND MLU_KERNEL_OBJECTS ${obj}) +endforeach() + +add_library(infini_train_mlu_kernels STATIC ${MLU_KERNEL_OBJECTS}) +set_target_properties(infini_train_mlu_kernels PROPERTIES LINKER_LANGUAGE CXX) +target_link_libraries(infini_train_mlu_kernels PUBLIC infini_train_mlu) + +# Static registrations (framework, CPU kernels, MLU runtime and MLU kernels) are +# only discovered through global initialization, so final executables must link +# them whole-archive. Reuse the framework helper instead of spelling out the +# link-group rules here. +add_library(infini_train_backend_mlu_executable INTERFACE) +add_library(InfiniTrain::Backend::MLUExecutable ALIAS + infini_train_backend_mlu_executable) +link_infini_train_exe(infini_train_backend_mlu_executable + EXTRA_ARCHIVES infini_train_mlu infini_train_mlu_kernels) + +if(BUILD_TEST) + message(FATAL_ERROR + "The MLU provider does not wire the shared InfiniTrain test suites yet. " + "Configure with -DBUILD_TEST=OFF and use the MLU regression tools " + "(-DINFINITRAIN_MLU_BUILD_TEST_TOOLS=ON) instead of reporting CPU test " + "results for the MLU device.") +endif() + +if(INFINITRAIN_MLU_BUILD_TEST_TOOLS) + add_subdirectory(tests) +endif() + +if(INFINITRAIN_BACKENDS_BUILD_EXAMPLES) + add_subdirectory(examples) +endif() diff --git a/backends/mlu/README.md b/backends/mlu/README.md new file mode 100644 index 0000000..f340c15 --- /dev/null +++ b/backends/mlu/README.md @@ -0,0 +1,98 @@ +# MLU backend (Cambricon) + +`backends/mlu` registers a `PrivateUse1` backend named `mlu` for Cambricon MLU devices. It is a +**work-in-progress** integration: see [`docs/validation.md`](docs/validation.md) for what has +actually been run and what has not. + +## Architecture in one paragraph + +InfiniTrain exposes third-party devices through the `PrivateUse1` mechanism: no new +`DeviceType` enum and no InfiniTrain source file changes. This provider registers the backend +from `src/backend.cc` (`infini_train::backends::mlu::RegisterBackend`): + +- `src/runtime/` - device guard, stream (CNRT queue), event (notifier), dispatch helpers; +- `src/kernels/*.mlu` - pure BANG C device code, compiled by `cncc`; +- `src/kernels/*.cc` - host-side integration (dispatch + `REGISTER_KERNEL`), compiled by the host compiler; +- `src/kernels/mlu_kernel_api.h` - the launcher interface between the two; +- `include/infini_train_mlu/backend.h` - the only public header: the registration entry point + used by examples and tools; +- `third_party/bang_host/` - two MIT-licensed Cambricon headers (`bang_fp16.h`, `bang_bf16.h`) + vendored for host compilation; see the README there for why and for the single local + modification. + +## Requirements + +| Item | Version used during development | +|---|---| +| Device | MLU590 (tested on MLU590-M9) | +| Driver | v6.5.52 | +| CNRT | 7.7.3 | +| cncc | 5.7.6 (clang 11.1 based) | +| Host compiler | GCC 13 (the framework needs C++20; GCC 11 lacks ``) | +| CMake / Ninja | CMake >= 3.28, Ninja | +| InfiniTrain | the pinned submodule (`v0.7.0`) | + +`CNNL` is optional at configure time: if `cnnl.h` / `libcnnl` are not found the build prints a +warning and disables BLAS handle support. **`INFINITRAIN_MLU_WITH_CNNL=ON` currently only +controls whether the CNNL handle is compiled in; no kernel calls CNNL yet** (GEMM still uses the +host fallback, see [`docs/validation.md`](docs/validation.md)). + +## Build + +```bash +git submodule update --init --recursive + +# CMAKE_POLICY_VERSION_MINIMUM=3.5 is required with CMake 4.x (third_party/gflags in the +# framework declares an older minimum). +cmake -S . -B build-mlu -G Ninja \ + -DINFINITRAIN_BACKEND=mlu -DUSE_CUDA=OFF -DBUILD_TEST=OFF \ + -DCMAKE_POLICY_VERSION_MINIMUM=3.5 +cmake --build build-mlu --target mlu_gpt2 mlu_op_test -j +``` + +Targets: `infini_train_mlu` (runtime + host kernel integration), `infini_train_mlu_kernels` +(BANG C objects), `mlu_gpt2` (upstream GPT-2 example with this backend registered, written as +`gpt2` in the build root) and `mlu_op_test` (operator regression tool, written as `op_test`). +`-DINFINITRAIN_MLU_BUILD_TEST_TOOLS=OFF` skips the regression tool; +`-DINFINITRAIN_BACKENDS_BUILD_EXAMPLES=OFF` skips the example. + +`BUILD_TEST=ON` is rejected on purpose: this provider does not wire the shared InfiniTrain test +suites yet, and CPU test results must not stand in for MLU validation. Use the MLU regression +tools instead. + +## Run + +```bash +GLOG_logtostderr=1 OMP_NUM_THREADS=64 ./build-mlu/gpt2 \ + --device mlu \ + --input_bin /tiny_shakespeare_train.bin \ + --llmc_filepath /gpt2_124M.bin \ + --num_iteration 1 +``` + +`GLOG_logtostderr=1` is required to see the step logs: glog otherwise writes them to +`/tmp/gpt2.*.INFO.*`. + +The command above is the configuration under validation - **not** a verified 10-step run. + +## Regression tools + +```bash +./build-mlu/op_test # operator vs CPU comparison (device test) +./build-mlu/op_test --inject-finite # negative check: the comparator must fail +./build-mlu/op_test --debug-status-test failure # controlled fatal path (no device fault) +cd backends/mlu/tests/math_selftest && bash build_and_run.sh # device math functions (this cwd) +``` + +`tests/math_selftest` has no framework dependency: it compiles its own `.mlu` with `cncc` and +the host part with `g++`, so it runs outside the CMake build. + +## Debug switches + +| Variable | Effect | +|---|---| +| `MLU_DEBUG_SYNC=1` | synchronize the queue after every kernel launch and log its name; an SDK failure aborts the process | +| `MLU_DEBUG_CHECK=1` | copy back a sample of every kernel output and log dtype / sampled count / non-finite count / min / max | + +Both parse `0`/`1`; unset or `0` disables them. They add synchronization and must be off for +any timing or acceptance run. diff --git a/backends/mlu/cmake/pre_project.cmake b/backends/mlu/cmake/pre_project.cmake new file mode 100644 index 0000000..fd7d97b --- /dev/null +++ b/backends/mlu/cmake/pre_project.cmake @@ -0,0 +1,30 @@ +# Cambricon MLU provider: settings that must be known before project(). +# +# The host compiler stays whatever the user selected (the framework needs C++20, +# GCC 13 in the validated environment). BANG C device sources are compiled by +# cncc through a custom command in this provider's CMakeLists, so no compiler +# replacement happens here. + +set(_MLU_DEFAULT_NEUWARE "$ENV{NEUWARE_HOME}") +if(NOT _MLU_DEFAULT_NEUWARE) + set(_MLU_DEFAULT_NEUWARE "/usr/local/neuware") +endif() +set(NEUWARE_HOME "${_MLU_DEFAULT_NEUWARE}" CACHE PATH + "Cambricon Neuware SDK root (CNRT, cncc and optional CNNL/CNCL)") +unset(_MLU_DEFAULT_NEUWARE) + +if(NOT EXISTS "${NEUWARE_HOME}/bin/cncc") + message(FATAL_ERROR + "The BANG C compiler was not found at ${NEUWARE_HOME}/bin/cncc. " + "Set NEUWARE_HOME or pass -DNEUWARE_HOME=.") +endif() +if(NOT EXISTS "${NEUWARE_HOME}/include/cnrt.h") + message(FATAL_ERROR + "CNRT headers were not found at ${NEUWARE_HOME}/include/cnrt.h.") +endif() + +# The framework is built for its CPU backend; CUDA/NCCL belong to other +# providers and multi-card collectives are not implemented for MLU yet. +set(USE_CUDA OFF) +set(USE_NCCL OFF) +set(BUILD_SHARED_LIBS OFF) diff --git a/backends/mlu/docs/validation.md b/backends/mlu/docs/validation.md new file mode 100644 index 0000000..2960dfe --- /dev/null +++ b/backends/mlu/docs/validation.md @@ -0,0 +1,64 @@ +# MLU backend - validation status (stage snapshot) + +This file records what has actually been run and what has not. It is deliberately conservative: +anything not listed as verified below should be treated as unverified. + +## Verified on 2026-11-03 (this provider layout) + +Built from a fresh CMake cache against the pinned InfiniTrain content (`c90338c`, v0.7.0) with +GCC 13.4.0, cncc 5.7.6, CMake 4.4.5 and Ninja 1.13.2; targets `mlu_gpt2` and `mlu_op_test`. + +| Item | Evidence | +|---|---| +| Fresh configuration and build | `configure` and `build` of `mlu_gpt2` + `mlu_op_test` both succeeded; all **9** BANG C sources were compiled by `cncc` | +| Operator comparison vs CPU (`op_test`) | 35 small-shape FP32 cases: Linear (both weight layouts, bias on/off, non-square 3x5->7), LayerNorm, Softmax, Add broadcast, MulScalar, Transpose, CrossEntropy, Pow domain and Pow(x,3) fwd/bwd. All pass with `abs(got-ref) <= 1e-5 + 1e-4*abs(ref)` (Pow uses 1e-6/1e-5) | +| Negative comparison paths | `--inject-finite` and `--inject-nan` return exit code 1 with `RESULT: FAIL`; the controlled fatal case aborts with the expected marker and no post-fatal work | +| Device math functions (`tests/math_selftest`) | exp / log / tanh / sqrt / rsqrt / pow / identity: 0/1024 mismatches under the tool's combined tolerance (`abs <= 1e-6 + 1e-5*abs(ref)`); the largest absolute error is 4.883e-4 on `exp` at x=8.5, where the reference is 4914.77 (relative 1.19e-7) | +| Example executable | `gpt2 --help` loads and prints usage. **No model training step was run** | +| Absence of stale framework paths | the generated compile database contains no path to the previously used `5a7be78` source tree | + +Note: Linear and Matmul currently execute through a **host fallback** - the operator tests above +prove that fallback path, not a device GEMM. + +## Verified earlier (historical, pre-migration layout) + +These were observed while the backend still lived under `mlu/` next to a separate framework tree. +They are recorded for traceability and were **not** re-executed for this layout change. + +| Item | Historical evidence | +|---|---| +| Backend registration | the `--device mlu` validator accepted the backend and kernels dispatched on `kPrivateUse1` | +| Build dependency tracking | touching the launcher header rebuilt all `.mlu` objects; a following build did nothing | + +## Not verified + +| Item | Status | +|---|---| +| Standard recursive-submodule download | the validation build used an archive of the pinned framework tree plus shallow retrieval of the four dependencies at their pinned commits. File content matched the pinned revisions, but the plain `git submodule update --init --recursive` procedure was **not** exercised | +| Cross-provider regression | no MACA or CUDA build was run. The untouched official files are not evidence of cross-platform regression coverage | +| Device GEMM (CNNL) | not implemented; `src/kernels/linear.cc` does D2H -> host GEMM -> H2D and is marked `NOTE(MLU-FALLBACK)` | +| Full GPT-2 FP32 single-card 10 steps | **not completed**; earlier attempts hit a queue failure at large batch and a slow single-core backward | +| Backward coverage at model scale | partial; small-shape backward cases exist, full-model backward has not finished | +| BF16 | registered in dispatch only; `src/kernels/embedding.mlu` still rounds per token (marked `TODO(BF16-MUST-FIX)`), no BF16 accuracy run | +| Multi-card (CNCL) | not implemented (communication kernels absent) | +| Shared InfiniTrain test suites | not wired; `BUILD_TEST=ON` is rejected on purpose and the MLU regression tools are used instead | +| Performance | kernels run single task / serial (`cnrtFuncTypeBlock`); no multi-core or NRAM vectorisation yet | + +## Known platform notes + +- `__bang_atomic_add` triggered a queue error in our usage, so `EmbeddingBackward` accumulates + with a row-partitioned scan instead (deterministic, no atomics). +- Device `powf` works once kernels run single-task/serial with `cnrtFuncTypeBlock`; with the + earlier multi-task `cnrtFuncTypeUnion1` launch, loop writes to global memory did not persist. +- `MluGuardImpl::Memcpy` must copy through the current stream (`cnrtMemcpyAsync` + queue sync); + plain `cnrtMemcpy` reads stale data while kernels are in flight. + +## Next steps + +1. device GEMM via CNNL (with the handle bound to the current stream) and re-run operator tests; +2. locate the large-shape queue failure, then run GPT-2 1 -> 2 -> 10 steps and compare loss + step by step against a frozen A100 baseline; +3. BF16: FP32 accumulation in embedding, autocast/dtype promotion, softmax/LayerNorm/CE stability; +4. multi-card: CNCL registration and a 2-card DP milestone; +5. wire the shared InfiniTrain test suites and run a MACA/CUDA regression before claiming + cross-provider compatibility. diff --git a/backends/mlu/examples/CMakeLists.txt b/backends/mlu/examples/CMakeLists.txt new file mode 100644 index 0000000..fb12e5d --- /dev/null +++ b/backends/mlu/examples/CMakeLists.txt @@ -0,0 +1,25 @@ +function(add_mlu_example target_name output_name) + add_executable(${target_name} ${ARGN}) + # Reuse the upstream example sources while injecting this provider's + # declaration and registration callback. + target_compile_definitions(${target_name} PRIVATE + INFINITRAIN_EXAMPLE_EXTERNAL_BACKEND_HEADER=\"infini_train_mlu/backend.h\" + INFINITRAIN_EXAMPLE_EXTERNAL_BACKEND_REGISTRAR=infini_train::backends::mlu::RegisterBackend + ) + target_link_libraries(${target_name} PRIVATE + InfiniTrain::Backend::MLUExecutable) + set_target_properties(${target_name} PROPERTIES + OUTPUT_NAME ${output_name} + RUNTIME_OUTPUT_DIRECTORY "${CMAKE_BINARY_DIR}" + ) +endfunction() + +# GPT-2 is the only model target for the current MLU stage. The sources come +# straight from the pinned InfiniTrain tree; nothing is copied or patched here. +add_mlu_example(mlu_gpt2 gpt2 + "${INFINITRAIN_SOURCE_DIR}/example/gpt2/main.cc" + "${INFINITRAIN_SOURCE_DIR}/example/gpt2/checkpoint_loader.cc" + "${INFINITRAIN_SOURCE_DIR}/example/common/tiny_shakespeare_dataset.cc" + "${INFINITRAIN_SOURCE_DIR}/example/common/tokenizer.cc" + "${INFINITRAIN_SOURCE_DIR}/example/common/utils.cc" +) diff --git a/backends/mlu/include/infini_train_mlu/backend.h b/backends/mlu/include/infini_train_mlu/backend.h new file mode 100644 index 0000000..9738101 --- /dev/null +++ b/backends/mlu/include/infini_train_mlu/backend.h @@ -0,0 +1,12 @@ +#pragma once + +// MLU（寒武纪）后端的注册入口。上游示例程序通过编译定义 +// INFINITRAIN_EXAMPLE_EXTERNAL_BACKEND_HEADER / _REGISTRAR 在解析命令行 +// 之前调用 RegisterBackend()（必须在 gflags 校验 --device 之前完成注册， +// 因为 Device::ParseType("mlu") 依赖注册后的 backend 名字）。 + +namespace infini_train::backends::mlu { + +void RegisterBackend(); + +} // namespace infini_train::backends::mlu diff --git a/backends/mlu/src/backend.cc b/backends/mlu/src/backend.cc new file mode 100644 index 0000000..59ffa16 --- /dev/null +++ b/backends/mlu/src/backend.cc @@ -0,0 +1,25 @@ +#include "infini_train_mlu/backend.h" + +#include + +#include "infini_train/include/core/privateuse1_backend.h" +#include "infini_train/include/datatype.h" + +namespace infini_train::backends::mlu { + +void RegisterBackend() { + static std::once_flag once; + std::call_once(once, [] { + // 前提（RegisterPrivateUse1Backend 内部 CHECK）：kPrivateUse1 的 + // DeviceGuardImpl 和 Cast/Fill/NoOpForward/NoOpBackward 四个必需 kernel + // 已通过静态注册就位——这依赖最终可执行文件以 --whole-archive 方式 + // 链接 infini_train_mlu / infini_train_mlu_kernels（由 provider 的 MLUExecutable 接口负责 whole-archive + // 链接）。 + core::PrivateUse1BackendRegistration registration; + registration.name = "mlu"; + registration.default_autocast_dtype = DataType::kBFLOAT16; + core::RegisterPrivateUse1Backend(registration); + }); +} + +} // namespace infini_train::backends::mlu diff --git a/backends/mlu/src/common/common_mlu.h b/backends/mlu/src/common/common_mlu.h new file mode 100644 index 0000000..ffcfd75 --- /dev/null +++ b/backends/mlu/src/common/common_mlu.h @@ -0,0 +1,48 @@ +#pragma once + +// 对应 InfiniTrain include/common/cuda/common_cuda.h。 +// NOTE(MLU-VERIFY): CNRT/CNNL/CNCL 的函数名以服务器上 Neuware 头文件 +// （${NEUWARE_HOME}/include/cnrt.h、cnnl.h、cncl.h）为准核一遍。 + +#include +#ifdef USE_CNNL +#include +#endif +#ifdef USE_CNCL +#include +#endif + +#include "glog/logging.h" + +namespace infini_train::common::mlu { + +#define CNRT_CHECK(call) \ + do { \ + cnrtRet_t status = (call); \ + if (status != cnrtSuccess) { \ + LOG(FATAL) << "CNRT Error: " << cnrtGetErrorStr(status) << " at " << __FILE__ << ":" << __LINE__; \ + } \ + } while (0) + +#ifdef USE_CNNL +#define CNNL_CHECK(call) \ + do { \ + cnnlStatus_t status = (call); \ + if (status != CNNL_STATUS_SUCCESS) { \ + LOG(FATAL) << "CNNL Error: " << cnnlGetErrorString(status) << " at " << __FILE__ << ":" << __LINE__; \ + } \ + } while (0) +#endif + +#ifdef USE_CNCL +#define CNCL_CHECK(expr) \ + do { \ + cnclResult_t _status = (expr); \ + if (_status != cnclSuccess) { \ + LOG(FATAL) << "CNCL error: " << cnclGetErrorStr(_status) << " at " << __FILE__ << ":" << __LINE__ << " (" \ + << #expr << ")"; \ + } \ + } while (0) +#endif + +} // namespace infini_train::common::mlu diff --git a/backends/mlu/src/kernels/accumulate_grad.cc b/backends/mlu/src/kernels/accumulate_grad.cc new file mode 100644 index 0000000..59d566b --- /dev/null +++ b/backends/mlu/src/kernels/accumulate_grad.cc @@ -0,0 +1,64 @@ +// 对应上游 src/kernels/cuda/accumulate_grad.cu 的 host 集成部分；设备 kernel 在 accumulate_grad.mlu。 + +#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 "kernels/mlu_kernel_api.h" +#include "runtime/mlu_dispatch.h" +#include "runtime/mlu_runtime_common.h" + +namespace infini_train::kernels::mlu { + +void AccumulateGrad(const std::shared_ptr &gradient, float rate, const std::shared_ptr &tensor) { + auto device = tensor->GetDevice(); + const auto queue = dynamic_cast( + infini_train::core::GetDeviceGuardImpl(device.type())->GetStream(device)) + ->mlu_queue(); + + core::mlu::DispatchMluFunc( + gradient->Dtype(), + [=]() { + AccumulateGradLaunch(gradient->DataPtr(), rate, tensor->DataPtr(), MluTypeTag(), + static_cast(gradient->NumElements()), queue); + ::infini_train::core::mlu::DebugSyncQueue(queue, "AccumulateGrad"); + }, + "MLU AccumulateGrad"); +} + +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) { + const float bias_correction_m = 1.0f - std::pow(beta1, t); + const float bias_correction_v = 1.0f - std::pow(beta2, t); + + auto device = grad->GetDevice(); + const auto queue = dynamic_cast( + infini_train::core::GetDeviceGuardImpl(device.type())->GetStream(device)) + ->mlu_queue(); + + core::mlu::DispatchMluFunc( + grad->Dtype(), + [=]() { + AdamAccumulateGradLaunch(grad->DataPtr(), param->DataPtr(), m->DataPtr(), v->DataPtr(), learning_rate, + beta1, beta2, eps, bias_correction_m, bias_correction_v, MluTypeTag(), + static_cast(grad->NumElements()), queue); + ::infini_train::core::mlu::DebugSyncQueue(queue, "AdamAccumulateGrad"); + }, + "MLU AdamAccumulateGrad"); +} +} // namespace infini_train::kernels::mlu + +#define REGISTER_MLU_ACCUMULATE_GRAD_KERNEL(kernel_name) \ + REGISTER_KERNEL(infini_train::Device::DeviceType::kPrivateUse1, kernel_name, \ + infini_train::kernels::mlu::kernel_name) + +REGISTER_MLU_ACCUMULATE_GRAD_KERNEL(AccumulateGrad) +REGISTER_MLU_ACCUMULATE_GRAD_KERNEL(AdamAccumulateGrad) + +#undef REGISTER_MLU_ACCUMULATE_GRAD_KERNEL diff --git a/backends/mlu/src/kernels/accumulate_grad.mlu b/backends/mlu/src/kernels/accumulate_grad.mlu new file mode 100644 index 0000000..79a156f --- /dev/null +++ b/backends/mlu/src/kernels/accumulate_grad.mlu @@ -0,0 +1,93 @@ +// 纯 BANG C kernel + 启动器（cncc 编译，只允许 C++11 级别语法，不含框架头）。 +// 框架集成见同目录 accumulate_grad.cc。 +// 语义对应上游 src/kernels/cuda/accumulate_grad.cu；16 位浮点一律经 float +// 中转计算（FP32 中间精度），与框架"归约/累加用 FP32"的约定一致。 + +#include + +#include + +#include "kernels/mlu_kernel_api.h" + +namespace { +constexpr uint32_t kDefaultTaskCount = 4; + +template +__mlu_global__ void AccumulateGradKernel(const T *grad_ptr, float rate, T *tensor_ptr, int64_t num_elements) { + + for (int64_t idx = 0; idx < num_elements; ++idx) { + tensor_ptr[idx] + = static_cast(static_cast(tensor_ptr[idx]) + static_cast(grad_ptr[idx]) * rate); + } +} + +template +__mlu_global__ void AdamAccumulateGradKernel(const T *grad_data, T *param_data, int64_t num_elements, T *m_data, + T *v_data, float learning_rate, float beta1, float beta2, float eps, + float bias_correction_m, float bias_correction_v) { + + for (int64_t idx = 0; idx < num_elements; ++idx) { + const float g = static_cast(grad_data[idx]); + const float m = beta1 * static_cast(m_data[idx]) + (1.0f - beta1) * g; + const float v = beta2 * static_cast(v_data[idx]) + (1.0f - beta2) * g * g; + m_data[idx] = static_cast(m); + v_data[idx] = static_cast(v); + + const float m_hat = m / bias_correction_m; + const float v_hat = v / bias_correction_v; + // CUDA 版用 __frcp_rn(__fsqrt_rn(v) + eps)（单精度倒数）；这里 1.0f/(sqrtf+eps) + // 同为 IEEE 单精度除法，对 FP32 场景逐位一致 + param_data[idx] + = static_cast(static_cast(param_data[idx]) - learning_rate * m_hat / (sqrtf(v_hat) + eps)); + } +} + +template +void AccumulateGradTyped(const void *grad, float rate, void *tensor, int64_t num_elements, cnrtQueue_t queue) { + const cnrtDim3_t dim = {1, 1, 1}; + AccumulateGradKernel + <<>>(static_cast(grad), rate, static_cast(tensor), num_elements); +} + +template +void AdamAccumulateGradTyped(const void *grad, void *param, void *m, void *v, float learning_rate, float beta1, + float beta2, float eps, float bias_correction_m, float bias_correction_v, + int64_t num_elements, cnrtQueue_t queue) { + const cnrtDim3_t dim = {1, 1, 1}; + AdamAccumulateGradKernel<<>>( + static_cast(grad), static_cast(param), num_elements, static_cast(m), static_cast(v), + learning_rate, beta1, beta2, eps, bias_correction_m, bias_correction_v); +} + +#define MLU_ACCUM_GRAD_FLOAT_SWITCH(launch_fn, ...) \ + switch (type) { \ + case MluScalarType::kF16: \ + launch_fn(__VA_ARGS__); \ + return; \ + case MluScalarType::kBF16: \ + launch_fn(__VA_ARGS__); \ + return; \ + case MluScalarType::kF32: \ + launch_fn(__VA_ARGS__); \ + return; \ + default: \ + /* 只 dispatch 浮点类型；FP64 已由 host 侧排除 */ \ + return; \ + } +} // namespace + +namespace infini_train::kernels::mlu { + +void AccumulateGradLaunch(const void *grad, float rate, void *tensor, MluScalarType type, int64_t num_elements, + cnrtQueue_t queue) { + MLU_ACCUM_GRAD_FLOAT_SWITCH(AccumulateGradTyped, grad, rate, tensor, num_elements, queue) +} + +void AdamAccumulateGradLaunch(const void *grad, void *param, void *m, void *v, float learning_rate, float beta1, + float beta2, float eps, float bias_correction_m, float bias_correction_v, + MluScalarType type, int64_t num_elements, cnrtQueue_t queue) { + MLU_ACCUM_GRAD_FLOAT_SWITCH(AdamAccumulateGradTyped, grad, param, m, v, learning_rate, beta1, beta2, eps, + bias_correction_m, bias_correction_v, num_elements, queue) +} + +} // namespace infini_train::kernels::mlu diff --git a/backends/mlu/src/kernels/cast.cc b/backends/mlu/src/kernels/cast.cc new file mode 100644 index 0000000..ce9230f --- /dev/null +++ b/backends/mlu/src/kernels/cast.cc @@ -0,0 +1,45 @@ +// 对应上游 src/kernels/cuda/cast.cu 的 host 集成部分；设备 kernel 在 cast.mlu。 + +#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/dtype_dispatch.h" +#include "infini_train/include/tensor.h" + +#include "kernels/mlu_kernel_api.h" +#include "runtime/mlu_dispatch.h" +#include "runtime/mlu_runtime_common.h" + +namespace infini_train::kernels::mlu { + +std::shared_ptr Cast(std::shared_ptr input, DataType dtype) { + auto device = input->GetDevice(); + auto dst_tensor = std::make_shared(input->Dims(), dtype, device); + const auto queue = dynamic_cast( + infini_train::core::GetDeviceGuardImpl(device.type())->GetStream(device)) + ->mlu_queue(); + + core::mlu::DispatchMluFunc, + DataTypeList>( + {dtype, input->Dtype()}, + [=]() { + CastLaunch(dst_tensor->DataPtr(), input->DataPtr(), MluTypeTag(), MluTypeTag(), + static_cast(input->NumElements()), queue); + ::infini_train::core::mlu::DebugSyncQueue(queue, "Cast"); + }, + "MLU Cast"); + + return {dst_tensor}; +} +} // namespace infini_train::kernels::mlu + +#define REGISTER_MLU_CAST_KERNEL(kernel_name) \ + REGISTER_KERNEL(infini_train::Device::DeviceType::kPrivateUse1, kernel_name, \ + infini_train::kernels::mlu::kernel_name) + +REGISTER_MLU_CAST_KERNEL(Cast) + +#undef REGISTER_MLU_CAST_KERNEL diff --git a/backends/mlu/src/kernels/cast.mlu b/backends/mlu/src/kernels/cast.mlu new file mode 100644 index 0000000..625c1af --- /dev/null +++ b/backends/mlu/src/kernels/cast.mlu @@ -0,0 +1,129 @@ +// 纯 BANG C kernel + 启动器（cncc 编译，只允许 C++11 级别语法，不含框架头）。 +// 框架集成见同目录 cast.cc。 +// NOTE(MLU-VERIFY): BANG C 中 half/bfloat16_t 与整型互转的 static_cast 行为 +// （截断/舍入语义）需与 baseline 对齐时复核。 + +#include +#include + +#include + +#include "kernels/mlu_kernel_api.h" + +namespace { +constexpr uint32_t kDefaultTaskCount = 4; + +template __mlu_global__ void CastKernel(Tdst *dst, const Tsrc *src, int64_t size) { + + for (int64_t idx = 0; idx < size; ++idx) { + constexpr bool src_reduced = std::is_same_v || std::is_same_v; + constexpr bool dst_reduced = std::is_same_v || std::is_same_v; + if constexpr (src_reduced || dst_reduced) { + // 任一侧是 16 位浮点都经 float 中转：BANG C 下 bf16 直接 static_cast + // 到整型/double 有歧义，且与框架 Cast"经 float 两步走"语义一致 + dst[idx] = static_cast(static_cast(src[idx])); + } else { + dst[idx] = static_cast(src[idx]); + } + } +} + +template +void CastTyped(void *dst, const void *src, int64_t num_elements, cnrtQueue_t queue) { + const cnrtDim3_t dim = {1, 1, 1}; + CastKernel + <<>>(static_cast(dst), static_cast(src), num_elements); +} + +#define MLU_CAST_SRC_SWITCH(Tdst) \ + switch (src_type) { \ + case MluScalarType::kBool: \ + CastTyped(dst, src, num_elements, queue); \ + return; \ + case MluScalarType::kU8: \ + CastTyped(dst, src, num_elements, queue); \ + return; \ + case MluScalarType::kI8: \ + CastTyped(dst, src, num_elements, queue); \ + return; \ + case MluScalarType::kU16: \ + CastTyped(dst, src, num_elements, queue); \ + return; \ + case MluScalarType::kI16: \ + CastTyped(dst, src, num_elements, queue); \ + return; \ + case MluScalarType::kU32: \ + CastTyped(dst, src, num_elements, queue); \ + return; \ + case MluScalarType::kI32: \ + CastTyped(dst, src, num_elements, queue); \ + return; \ + case MluScalarType::kU64: \ + CastTyped(dst, src, num_elements, queue); \ + return; \ + case MluScalarType::kI64: \ + CastTyped(dst, src, num_elements, queue); \ + return; \ + case MluScalarType::kF16: \ + CastTyped(dst, src, num_elements, queue); \ + return; \ + case MluScalarType::kBF16: \ + CastTyped(dst, src, num_elements, queue); \ + return; \ + case MluScalarType::kF32: \ + CastTyped(dst, src, num_elements, queue); \ + return; \ + default: \ + /* MLU 设备不支持 FP64（kF64 不落到这里；host 侧 dispatch 已排除）*/ \ + return; \ + } +} // namespace + +namespace infini_train::kernels::mlu { + +void CastLaunch(void *dst, const void *src, MluScalarType dst_type, MluScalarType src_type, int64_t num_elements, + cnrtQueue_t queue) { + switch (dst_type) { + case MluScalarType::kBool: + MLU_CAST_SRC_SWITCH(bool) + return; + case MluScalarType::kU8: + MLU_CAST_SRC_SWITCH(uint8_t) + return; + case MluScalarType::kI8: + MLU_CAST_SRC_SWITCH(int8_t) + return; + case MluScalarType::kU16: + MLU_CAST_SRC_SWITCH(uint16_t) + return; + case MluScalarType::kI16: + MLU_CAST_SRC_SWITCH(int16_t) + return; + case MluScalarType::kU32: + MLU_CAST_SRC_SWITCH(uint32_t) + return; + case MluScalarType::kI32: + MLU_CAST_SRC_SWITCH(int32_t) + return; + case MluScalarType::kU64: + MLU_CAST_SRC_SWITCH(uint64_t) + return; + case MluScalarType::kI64: + MLU_CAST_SRC_SWITCH(int64_t) + return; + case MluScalarType::kF16: + MLU_CAST_SRC_SWITCH(half) + return; + case MluScalarType::kBF16: + MLU_CAST_SRC_SWITCH(bfloat16_t) + return; + case MluScalarType::kF32: + MLU_CAST_SRC_SWITCH(float) + return; + default: + // MLU 设备不支持 FP64（kF64 已由 host 侧 dispatch 排除） + return; + } +} + +} // namespace infini_train::kernels::mlu diff --git a/backends/mlu/src/kernels/cross_entropy.cc b/backends/mlu/src/kernels/cross_entropy.cc new file mode 100644 index 0000000..26098dc --- /dev/null +++ b/backends/mlu/src/kernels/cross_entropy.cc @@ -0,0 +1,101 @@ +// 对应上游 src/kernels/cuda/cross_entropy.cu 的 host 集成部分；设备 kernel 在 cross_entropy.mlu。 +// 与 CUDA 版一致：kernel 只算每样本 loss，批均值在 host 侧 FP64 累加（含 D2H/H2D）。 + +#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/device.h" +#include "infini_train/include/dispatcher.h" +#include "infini_train/include/tensor.h" + +#include "kernels/mlu_kernel_api.h" +#include "runtime/mlu_debug.h" +#include "runtime/mlu_dispatch.h" +#include "runtime/mlu_runtime_common.h" + +namespace infini_train::kernels::mlu { +namespace { +cnrtQueue_t CurrentQueue(const Device &device) { + return dynamic_cast( + infini_train::core::GetDeviceGuardImpl(device.type())->GetStream(device)) + ->mlu_queue(); +} +} // namespace + +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 int64_t bs = std::accumulate(input_dims.rbegin() + 1, input_dims.rend(), 1, std::multiplies{}); + const int64_t num_classes = *input_dims.rbegin(); + + auto batched_output = std::make_shared(std::vector{bs}, input->Dtype(), input->GetDevice()); + const auto queue = CurrentQueue(input->GetDevice()); + + return core::mlu::DispatchMluFunc, + DataTypeList>( + {target->Dtype(), input->Dtype()}, + [=]() { + CrossEntropyForwardLaunch(input->DataPtr(), target->DataPtr(), batched_output->DataPtr(), bs, num_classes, + MluTypeTag(), MluTypeTag(), queue); + // 探针放在 Launch 之后、第一次 To(Device())（含 D2H）之前； + // 原先写在 return 之后属死代码，导致这个最重的 kernel 从未被覆盖。 + ::infini_train::core::mlu::DebugSyncQueue(queue, "CrossEntropyForward"); + ::infini_train::core::mlu::DebugCheck(batched_output, "CrossEntropyForward"); + + // 与 CUDA 版一致：host 侧 FP64 累加求均值（含 D2H/H2D 同步） + auto loss_cpu = batched_output->To(Device()); + auto loss = std::make_shared(std::vector{}, input->Dtype(), Device()); + const auto *loss_cpu_ptr = static_cast(loss_cpu.DataPtr()); + const double loss_sum + = std::accumulate(loss_cpu_ptr, loss_cpu_ptr + bs, 0.0, + [](double acc, const Tinput &val) { return acc + static_cast(val); }); + const float loss_mean = static_cast(loss_sum / static_cast(bs)); + static_cast(loss->DataPtr())[0] = static_cast(loss_mean); + + return std::make_shared(loss->To(input->GetDevice())); + }, + "MLU CrossEntropyForward"); +} + +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 int64_t bs = std::accumulate(input_dims.rbegin() + 1, input_dims.rend(), 1, std::multiplies{}); + const int64_t 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()); + + const auto queue = CurrentQueue(target->GetDevice()); + core::mlu::DispatchMluFunc, + DataTypeList>( + {target->Dtype(), input_casted->Dtype()}, + [=]() { + CrossEntropyBackwardLaunch(input_casted->DataPtr(), grad_input->DataPtr(), target->DataPtr(), + grad_output->DataPtr(), bs, num_classes, MluTypeTag(), + MluTypeTag(), queue); + ::infini_train::core::mlu::DebugSyncQueue(queue, "CrossEntropyBackward"); + }, + "MLU CrossEntropyBackward"); + return {grad_input}; +} +} // namespace infini_train::kernels::mlu + +#define REGISTER_MLU_CROSS_ENTROPY_KERNEL(kernel_name) \ + REGISTER_KERNEL(infini_train::Device::DeviceType::kPrivateUse1, kernel_name, \ + infini_train::kernels::mlu::kernel_name) + +REGISTER_MLU_CROSS_ENTROPY_KERNEL(CrossEntropyForward) +REGISTER_MLU_CROSS_ENTROPY_KERNEL(CrossEntropyBackward) + +#undef REGISTER_MLU_CROSS_ENTROPY_KERNEL diff --git a/backends/mlu/src/kernels/cross_entropy.mlu b/backends/mlu/src/kernels/cross_entropy.mlu new file mode 100644 index 0000000..80636b3 --- /dev/null +++ b/backends/mlu/src/kernels/cross_entropy.mlu @@ -0,0 +1,130 @@ +// 纯 BANG C kernel + 启动器（cncc 编译，只允许 C++11 级别语法，不含框架头）。 +// 框架集成见同目录 cross_entropy.cc。 +// 语义对应上游 src/kernels/cuda/cross_entropy.cu：max/sum 归约全程 FP32。 +// 简化策略（正确性优先）：一个样本一个核，串行归约，顺序确定。 +// 注意：CUDA 版前向的批均值是在 host 侧 FP64 累加的（见其 FIXME 注释）， +// 本 kernel 只产出每样本 loss，批均值同样在 host 侧做（见 cross_entropy.cc）。 + +#include + +#include + +#include "kernels/mlu_kernel_api.h" + +namespace { +using infini_train::kernels::mlu::MluScalarType; + +constexpr uint32_t kDefaultTaskCount = 4; + +// 注意（实测）：本工具链下 taskIdX/coreId/taskDimX/coreDim 与预期语义不符，配合 +// cnrtFuncTypeUnion1 时循环写 GDRAM 只有首尾元素落盘。厂商示例（samples/BANG/.../ +// vectoradd/0_scalarAdd_single_core）用的是纯串行循环 + cnrtFuncTypeBlock。 +// 因此所有 kernel 统一改为单任务串行；性能优化属任务三，届时再按官方多核范式改写。 + +template +__mlu_global__ void CrossEntropyForwardKernel(const Tinput *input, const Ttarget *target, Tinput *loss, int64_t bs, + int64_t num_classes) { + for (int64_t sample = 0; sample < bs; ++sample) { + const Tinput *row = input + sample * num_classes; + + float max_logit = -__builtin_inff(); + for (int64_t i = 0; i < num_classes; ++i) { + const float v = static_cast(row[i]); + max_logit = v > max_logit ? v : max_logit; + } + float sum_exp = 0.0f; + for (int64_t i = 0; i < num_classes; ++i) { sum_exp += expf(static_cast(row[i]) - max_logit); } + + const int64_t target_class = static_cast(target[sample]); + const float target_val = static_cast(row[target_class]) - max_logit; + loss[sample] = static_cast(logf(sum_exp) - target_val); + } +} + +template +__mlu_global__ void CrossEntropyBackwardKernel(const Tinput *input, Tinput *grad_input, const Ttarget *target, + const Tinput *grad_output, int64_t bs, int64_t num_classes) { + const float inv_bs = 1.0f / bs; + const float grad_scale = static_cast(grad_output[0]); + for (int64_t sample = 0; sample < bs; ++sample) { + const Tinput *row = input + sample * num_classes; + Tinput *grad_row = grad_input + sample * num_classes; + + float max_logit = -__builtin_inff(); + for (int64_t i = 0; i < num_classes; ++i) { + const float v = static_cast(row[i]); + max_logit = v > max_logit ? v : max_logit; + } + float sum_exp = 0.0f; + for (int64_t i = 0; i < num_classes; ++i) { sum_exp += expf(static_cast(row[i]) - max_logit); } + + const int64_t target_class = static_cast(target[sample]); + const float scale = 1.0f / sum_exp; + for (int64_t i = 0; i < num_classes; ++i) { + const float exp_val = expf(static_cast(row[i]) - max_logit); + grad_row[i] + = static_cast((exp_val * scale - (i == target_class ? 1.0f : 0.0f)) * inv_bs * grad_scale); + } + } +} + +template +void CrossEntropyForwardTyped(const void *input, const void *target, void *loss, int64_t bs, int64_t num_classes, + cnrtQueue_t queue) { + const cnrtDim3_t dim = {1, 1, 1}; + CrossEntropyForwardKernel + <<>>(static_cast(input), static_cast(target), + static_cast(loss), bs, num_classes); +} + +template +void CrossEntropyBackwardTyped(const void *input, void *grad_input, const void *target, const void *grad_output, + int64_t bs, int64_t num_classes, cnrtQueue_t queue) { + const cnrtDim3_t dim = {1, 1, 1}; + CrossEntropyBackwardKernel<<>>( + static_cast(input), static_cast(grad_input), static_cast(target), + static_cast(grad_output), bs, num_classes); +} + +#define MLU_CE_INPUT_SWITCH(fn, Ttarget, ...) \ + switch (input_type) { \ + case MluScalarType::kF16: \ + fn(__VA_ARGS__); \ + return; \ + case MluScalarType::kBF16: \ + fn(__VA_ARGS__); \ + return; \ + case MluScalarType::kF32: \ + fn(__VA_ARGS__); \ + return; \ + default: \ + return; \ + } + +#define MLU_CE_TARGET_SWITCH(fn, ...) \ + switch (target_type) { \ + case MluScalarType::kU8: \ + MLU_CE_INPUT_SWITCH(fn, uint8_t, __VA_ARGS__); \ + return; \ + case MluScalarType::kI64: \ + MLU_CE_INPUT_SWITCH(fn, int64_t, __VA_ARGS__); \ + return; \ + default: \ + return; \ + } +} // namespace + +namespace infini_train::kernels::mlu { + +void CrossEntropyForwardLaunch(const void *input, const void *target, void *loss, int64_t bs, int64_t num_classes, + MluScalarType target_type, MluScalarType input_type, cnrtQueue_t queue) { + MLU_CE_TARGET_SWITCH(CrossEntropyForwardTyped, input, target, loss, bs, num_classes, queue) +} + +void CrossEntropyBackwardLaunch(const void *input, void *grad_input, const void *target, const void *grad_output, + int64_t bs, int64_t num_classes, MluScalarType target_type, MluScalarType input_type, + cnrtQueue_t queue) { + MLU_CE_TARGET_SWITCH(CrossEntropyBackwardTyped, input, grad_input, target, grad_output, bs, num_classes, queue) +} + +} // namespace infini_train::kernels::mlu diff --git a/backends/mlu/src/kernels/elementwise.cc b/backends/mlu/src/kernels/elementwise.cc new file mode 100644 index 0000000..84088fa --- /dev/null +++ b/backends/mlu/src/kernels/elementwise.cc @@ -0,0 +1,428 @@ +// 对应上游 src/kernels/cuda/elementwise.cu 的 host 集成部分；设备 kernel 在 elementwise.mlu。 +// 只保留 GPT-2 训练路径需要的算子（一元 Neg/Reciprocal/Sin/Cos/Tanh/Pow/Rsqrt/Exp/Log/Sigmoid、 +// 二元 Add/Sub/Mul/Div 及 Add/Mul 标量变体、Equals/EqualsScalar）。 + +#include +#include +#include +#include +#include +#include +#include + +#include "glog/logging.h" + +#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/scalar.h" +#include "infini_train/include/tensor.h" + +#include "kernels/mlu_kernel_api.h" +#include "runtime/mlu_debug.h" +#include "runtime/mlu_dispatch.h" +#include "runtime/mlu_runtime_common.h" + +namespace infini_train::kernels::mlu { +namespace { + +cnrtQueue_t CurrentQueue(const Device &device) { + return dynamic_cast( + infini_train::core::GetDeviceGuardImpl(device.type())->GetStream(device)) + ->mlu_queue(); +} + +// 对应 CUDA 版 MakeBroadcastMeta：a/b 的 shape 右对齐到 out 的 ndim +MluBroadcastMeta MakeBroadcastMeta(const std::vector &a_dims, const std::vector &b_dims, + const std::vector &out_dims) { + MluBroadcastMeta m{}; + const int ndim = static_cast(out_dims.size()); + CHECK_LE(ndim, MLU_MAX_BROADCAST_DIMS) << "Broadcast ndim exceeds MLU_MAX_BROADCAST_DIMS"; + 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()); + + const auto a_str = ComputeStrides(a_shape); + const auto b_str = ComputeStrides(b_shape); + const 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; +} + +bool ShapesEqual(const std::vector &a, const std::vector &b) { return a == b; } + +// dtype 提升（与 CUDA 版一致：按 PromoteDataTypes 结果 Cast） +std::shared_ptr Promote(const std::shared_ptr &t, DataType promoted) { + return t->Dtype() == promoted ? t : std::make_shared(t->To(promoted)); +} + +std::shared_ptr UnaryForwardImpl(const std::shared_ptr &input, MluUnaryOp op) { + auto output = std::make_shared(input->Dims(), input->Dtype(), input->GetDevice()); + const auto queue = CurrentQueue(input->GetDevice()); + core::mlu::DispatchMluFunc( + input->Dtype(), + [=]() { + UnaryForwardLaunch(output->DataPtr(), input->DataPtr(), op, MluTypeTag(), + static_cast(input->NumElements()), queue); + ::infini_train::core::mlu::DebugSyncQueue(queue, "UnaryForward"); + ::infini_train::core::mlu::DebugCheck(output, "UnaryForward"); + }, + "MLU UnaryForward"); + return output; +} + +// aux 语义见 mlu_kernel_api.h UnaryBackwardLaunch 注释 +std::shared_ptr UnaryBackwardImpl(const std::shared_ptr &grad_output, + const std::shared_ptr &aux, MluUnaryOp op) { + const auto dtype = grad_output->Dtype(); + const auto aux_dtype = aux ? aux->Dtype() : dtype; + const DataType promoted = PromoteDataTypes(dtype, aux_dtype); + auto g = Promote(grad_output, promoted); + auto x = aux ? Promote(aux, promoted) : nullptr; + auto grad_in = std::make_shared(grad_output->Dims(), promoted, grad_output->GetDevice()); + const auto queue = CurrentQueue(grad_output->GetDevice()); + core::mlu::DispatchMluFunc( + promoted, + [=]() { + UnaryBackwardLaunch(grad_in->DataPtr(), g->DataPtr(), x ? x->DataPtr() : g->DataPtr(), op, MluTypeTag(), + static_cast(g->NumElements()), queue); + ::infini_train::core::mlu::DebugSyncQueue(queue, "UnaryBackward"); + }, + "MLU UnaryBackward"); + return grad_in; +} + +std::shared_ptr BinaryForwardImpl(const std::shared_ptr &a, const std::shared_ptr &b, + MluBinaryOp op) { + const DataType promoted = PromoteDataTypes(a->Dtype(), b->Dtype()); + auto ap = Promote(a, promoted); + auto bp = Promote(b, promoted); + // 与 CUDA 版一致：只假设 b -> a 单向广播 + CHECK(a->NumElements() >= b->NumElements() && a->NumElements() % b->NumElements() == 0); + + auto output = std::make_shared(a->Dims(), promoted, a->GetDevice()); + const auto meta = MakeBroadcastMeta(ap->Dims(), bp->Dims(), output->Dims()); + const auto queue = CurrentQueue(a->GetDevice()); + core::mlu::DispatchMluFunc( + promoted, + [=]() { + BinaryForwardLaunch(output->DataPtr(), ap->DataPtr(), bp->DataPtr(), op, MluTypeTag(), meta, + static_cast(output->NumElements()), queue); + ::infini_train::core::mlu::DebugSyncQueue(queue, "BinaryForward"); + ::infini_train::core::mlu::DebugCheck(output, "BinaryForward"); + }, + "MLU BinaryForward"); + return output; +} + +std::pair, std::shared_ptr> +BinaryBackwardImpl(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, MluBinaryOp op) { + auto dtype = grad_output->Dtype(); + const auto a_dtype = a ? a->Dtype() : dtype; + const auto b_dtype = b ? b->Dtype() : dtype; + // 与 CUDA 版一致：计算 dtype 由前向保存的张量决定，而不是 grad_output + const DataType promoted = PromoteDataTypes(a_dtype, b_dtype); + auto g = Promote(grad_output, promoted); + auto ap = a ? Promote(a, promoted) : nullptr; + auto bp = b ? Promote(b, promoted) : nullptr; + // Add/Sub 的反向不需要 a/b 的值（grad_a=g, grad_b=±g），调用方传 nullptr； + // 用 g 占位保证 kernel 内索引不越界（公式本身不会读到） + if (!ap) { + ap = g; + } + if (!bp) { + bp = g; + } + + auto device = grad_output->GetDevice(); + auto grad_a = std::make_shared(a_dims, promoted, device); + auto grad_b = std::make_shared(b_dims, promoted, device); + + const auto meta = MakeBroadcastMeta(a_dims, b_dims, g->Dims()); + const auto queue = CurrentQueue(device); + const int64_t numel_a = grad_a->NumElements(); + const int64_t numel_b = grad_b->NumElements(); + const int64_t numel_out = g->NumElements(); + core::mlu::DispatchMluFunc( + promoted, + [=]() { + BinaryBackwardLaunch(grad_a->DataPtr(), grad_b->DataPtr(), g->DataPtr(), ap->DataPtr(), bp->DataPtr(), op, + MluTypeTag(), meta, numel_a, numel_b, numel_out, queue); + ::infini_train::core::mlu::DebugSyncQueue(queue, "BinaryBackward"); + }, + "MLU BinaryBackward"); + return {grad_a, grad_b}; +} + +std::shared_ptr ScalarForwardImpl(const std::shared_ptr &input, float value, MluBinaryOp op) { + auto output = std::make_shared(input->Dims(), input->Dtype(), input->GetDevice()); + const auto queue = CurrentQueue(input->GetDevice()); + core::mlu::DispatchMluFunc( + input->Dtype(), + [=]() { + ScalarForwardLaunch(output->DataPtr(), input->DataPtr(), value, op, MluTypeTag(), + static_cast(input->NumElements()), queue); + ::infini_train::core::mlu::DebugSyncQueue(queue, "ScalarForward"); + if (std::getenv("MLU_DEBUG_CHECK")) { + char tag[96]; + std::snprintf(tag, sizeof(tag), "ScalarForward(op=%d,scalar=%g) in", static_cast(op), value); + ::infini_train::core::mlu::DebugCheck(input, tag); + std::snprintf(tag, sizeof(tag), "ScalarForward(op=%d,scalar=%g) out", static_cast(op), value); + ::infini_train::core::mlu::DebugCheck(output, tag); + } + }, + "MLU ScalarForward"); + return output; +} + +std::shared_ptr ScalarBackwardImpl(const std::shared_ptr &grad_output, float value, MluBinaryOp op) { + auto grad_in = std::make_shared(grad_output->Dims(), grad_output->Dtype(), grad_output->GetDevice()); + const auto queue = CurrentQueue(grad_output->GetDevice()); + core::mlu::DispatchMluFunc( + grad_output->Dtype(), + [=]() { + ScalarBackwardLaunch(grad_in->DataPtr(), grad_output->DataPtr(), value, op, MluTypeTag(), + static_cast(grad_output->NumElements()), queue); + ::infini_train::core::mlu::DebugSyncQueue(queue, "ScalarBackward"); + }, + "MLU ScalarBackward"); + return grad_in; +} + +} // namespace + +// ----------------------------- 一元 ----------------------------- + +std::shared_ptr NegForward(const std::shared_ptr &input) { + return UnaryForwardImpl(input, MluUnaryOp::kNeg); +} +std::shared_ptr NegBackward(const std::shared_ptr &grad_output) { + return UnaryBackwardImpl(grad_output, nullptr, MluUnaryOp::kNeg); +} +std::shared_ptr ReciprocalForward(const std::shared_ptr &input) { + return UnaryForwardImpl(input, MluUnaryOp::kReciprocal); +} +std::shared_ptr ReciprocalBackward(const std::shared_ptr &grad_output, + const std::shared_ptr &input) { + return UnaryBackwardImpl(grad_output, input, MluUnaryOp::kReciprocal); +} +std::shared_ptr SinForward(const std::shared_ptr &input) { + return UnaryForwardImpl(input, MluUnaryOp::kSin); +} +std::shared_ptr SinBackward(const std::shared_ptr &grad_output, const std::shared_ptr &input) { + return UnaryBackwardImpl(grad_output, input, MluUnaryOp::kSin); +} +std::shared_ptr CosForward(const std::shared_ptr &input) { + return UnaryForwardImpl(input, MluUnaryOp::kCos); +} +std::shared_ptr CosBackward(const std::shared_ptr &grad_output, const std::shared_ptr &input) { + return UnaryBackwardImpl(grad_output, input, MluUnaryOp::kCos); +} +std::shared_ptr TanhForward(const std::shared_ptr &input) { + return UnaryForwardImpl(input, MluUnaryOp::kTanh); +} +std::shared_ptr TanhBackward(const std::shared_ptr &grad_output, + const std::shared_ptr &output) { + return UnaryBackwardImpl(grad_output, output, MluUnaryOp::kTanh); +} +std::shared_ptr RsqrtForward(const std::shared_ptr &input) { + return UnaryForwardImpl(input, MluUnaryOp::kRsqrt); +} +std::shared_ptr RsqrtBackward(const std::shared_ptr &grad_output, + const std::shared_ptr &input) { + return UnaryBackwardImpl(grad_output, input, MluUnaryOp::kRsqrt); +} +std::shared_ptr ExpForward(const std::shared_ptr &input) { + return UnaryForwardImpl(input, MluUnaryOp::kExp); +} +std::shared_ptr ExpBackward(const std::shared_ptr &grad_output, const std::shared_ptr &output) { + return UnaryBackwardImpl(grad_output, output, MluUnaryOp::kExp); +} +std::shared_ptr LogForward(const std::shared_ptr &input) { + return UnaryForwardImpl(input, MluUnaryOp::kLog); +} +std::shared_ptr LogBackward(const std::shared_ptr &grad_output, const std::shared_ptr &input) { + return UnaryBackwardImpl(grad_output, input, MluUnaryOp::kLog); +} +std::shared_ptr SigmoidForward(const std::shared_ptr &input) { + return UnaryForwardImpl(input, MluUnaryOp::kSigmoid); +} +// 注意调用方参数顺序是 (output, grad_output) +std::shared_ptr SigmoidBackward(const std::shared_ptr &output, + const std::shared_ptr &grad_output) { + return UnaryBackwardImpl(grad_output, output, MluUnaryOp::kSigmoid); +} + +std::shared_ptr PowForward(const std::shared_ptr &input, float exponent, bool scalar_is_base) { + auto output = std::make_shared(input->Dims(), input->Dtype(), input->GetDevice()); + const float e = exponent; + const auto queue = CurrentQueue(input->GetDevice()); + core::mlu::DispatchMluFunc( + input->Dtype(), + [=]() { + PowForwardLaunch(output->DataPtr(), input->DataPtr(), e, scalar_is_base, MluTypeTag(), + static_cast(input->NumElements()), queue); + ::infini_train::core::mlu::DebugSyncQueue(queue, "PowForward"); + ::infini_train::core::mlu::DebugCheck(output, "PowForward"); + }, + "MLU PowForward"); + return output; +} + +std::shared_ptr PowBackward(const std::shared_ptr &grad_output, const std::shared_ptr &input, + float exponent, bool scalar_is_base) { + const DataType promoted = PromoteDataTypes(grad_output->Dtype(), input->Dtype()); + auto g = Promote(grad_output, promoted); + auto x = Promote(input, promoted); + auto grad_in = std::make_shared(g->Dims(), promoted, g->GetDevice()); + const float e = exponent; + const auto queue = CurrentQueue(g->GetDevice()); + core::mlu::DispatchMluFunc( + promoted, + [=]() { + PowBackwardLaunch(grad_in->DataPtr(), g->DataPtr(), x->DataPtr(), e, scalar_is_base, MluTypeTag(), + static_cast(g->NumElements()), queue); + ::infini_train::core::mlu::DebugSyncQueue(queue, "PowBackward"); + }, + "MLU PowBackward"); + return grad_in; +} + +// ----------------------------- 比较 ----------------------------- + +std::shared_ptr EqualsForward(const std::shared_ptr &a, const std::shared_ptr &b) { + CHECK(a->Dtype() == b->Dtype()); + auto output = std::make_shared(a->Dims(), DataType::kBOOL, a->GetDevice()); + const auto meta = MakeBroadcastMeta(a->Dims(), b->Dims(), output->Dims()); + const auto queue = CurrentQueue(a->GetDevice()); + core::mlu::DispatchMluFunc( + a->Dtype(), + [=]() { + EqualsForwardLaunch(output->DataPtr(), a->DataPtr(), b->DataPtr(), MluTypeTag(), meta, + static_cast(output->NumElements()), queue); + ::infini_train::core::mlu::DebugSyncQueue(queue, "EqualsForward"); + }, + "MLU EqualsForward"); + return output; +} + +std::shared_ptr EqualsScalarForward(const std::shared_ptr &input, float scalar) { + auto output = std::make_shared(input->Dims(), DataType::kBOOL, input->GetDevice()); + const double value = static_cast(scalar); + const auto queue = CurrentQueue(input->GetDevice()); + core::mlu::DispatchMluFunc( + input->Dtype(), + [=]() { + EqualsScalarForwardLaunch(output->DataPtr(), input->DataPtr(), value, MluTypeTag(), + static_cast(input->NumElements()), queue); + ::infini_train::core::mlu::DebugSyncQueue(queue, "EqualsScalarForward"); + }, + "MLU EqualsScalarForward"); + return output; +} + +// ----------------------------- 二元 ----------------------------- + +std::shared_ptr AddForward(const std::shared_ptr &a, const std::shared_ptr &b) { + return BinaryForwardImpl(a, b, MluBinaryOp::kAdd); +} +std::pair, std::shared_ptr> AddBackward(const std::shared_ptr &grad_output, + const std::vector &a_dims, + const std::vector &b_dims) { + return BinaryBackwardImpl(grad_output, nullptr, nullptr, a_dims, b_dims, MluBinaryOp::kAdd); +} +std::shared_ptr AddScalarForward(const std::shared_ptr &input, float scalar) { + return ScalarForwardImpl(input, scalar, MluBinaryOp::kAdd); +} +std::shared_ptr AddScalarBackward(const std::shared_ptr &grad_output) { + return ScalarBackwardImpl(grad_output, 0.0f, MluBinaryOp::kAdd); +} + +std::shared_ptr SubForward(const std::shared_ptr &a, const std::shared_ptr &b) { + return BinaryForwardImpl(a, b, MluBinaryOp::kSub); +} +std::pair, std::shared_ptr> SubBackward(const std::shared_ptr &grad_output, + const std::vector &a_dims, + const std::vector &b_dims) { + return BinaryBackwardImpl(grad_output, nullptr, nullptr, a_dims, b_dims, MluBinaryOp::kSub); +} + +std::shared_ptr MulForward(const std::shared_ptr &a, const std::shared_ptr &b) { + return BinaryForwardImpl(a, b, MluBinaryOp::kMul); +} +std::pair, std::shared_ptr> MulBackward(const std::shared_ptr &grad_output, + const std::shared_ptr &a, + const std::shared_ptr &b) { + return BinaryBackwardImpl(grad_output, a, b, a->Dims(), b->Dims(), MluBinaryOp::kMul); +} +std::shared_ptr MulScalarForward(const std::shared_ptr &input, float scalar) { + return ScalarForwardImpl(input, scalar, MluBinaryOp::kMul); +} +std::shared_ptr MulScalarBackward(const std::shared_ptr &grad_output, float scalar) { + return ScalarBackwardImpl(grad_output, scalar, MluBinaryOp::kMul); +} + +std::shared_ptr DivForward(const std::shared_ptr &a, const std::shared_ptr &b) { + return BinaryForwardImpl(a, b, MluBinaryOp::kDiv); +} +std::pair, std::shared_ptr> DivBackward(const std::shared_ptr &grad_output, + const std::shared_ptr &a, + const std::shared_ptr &b) { + return BinaryBackwardImpl(grad_output, a, b, a->Dims(), b->Dims(), MluBinaryOp::kDiv); +} + +} // namespace infini_train::kernels::mlu + +#define REGISTER_MLU_ELEMENTWISE_KERNEL(kernel_name) \ + REGISTER_KERNEL(infini_train::Device::DeviceType::kPrivateUse1, kernel_name, \ + infini_train::kernels::mlu::kernel_name) + +REGISTER_MLU_ELEMENTWISE_KERNEL(NegForward) +REGISTER_MLU_ELEMENTWISE_KERNEL(NegBackward) +REGISTER_MLU_ELEMENTWISE_KERNEL(ReciprocalForward) +REGISTER_MLU_ELEMENTWISE_KERNEL(ReciprocalBackward) +REGISTER_MLU_ELEMENTWISE_KERNEL(SinForward) +REGISTER_MLU_ELEMENTWISE_KERNEL(SinBackward) +REGISTER_MLU_ELEMENTWISE_KERNEL(CosForward) +REGISTER_MLU_ELEMENTWISE_KERNEL(CosBackward) +REGISTER_MLU_ELEMENTWISE_KERNEL(TanhForward) +REGISTER_MLU_ELEMENTWISE_KERNEL(TanhBackward) +REGISTER_MLU_ELEMENTWISE_KERNEL(PowForward) +REGISTER_MLU_ELEMENTWISE_KERNEL(PowBackward) +REGISTER_MLU_ELEMENTWISE_KERNEL(RsqrtForward) +REGISTER_MLU_ELEMENTWISE_KERNEL(RsqrtBackward) +REGISTER_MLU_ELEMENTWISE_KERNEL(ExpForward) +REGISTER_MLU_ELEMENTWISE_KERNEL(ExpBackward) +REGISTER_MLU_ELEMENTWISE_KERNEL(LogForward) +REGISTER_MLU_ELEMENTWISE_KERNEL(LogBackward) +REGISTER_MLU_ELEMENTWISE_KERNEL(SigmoidForward) +REGISTER_MLU_ELEMENTWISE_KERNEL(SigmoidBackward) +REGISTER_MLU_ELEMENTWISE_KERNEL(EqualsForward) +REGISTER_MLU_ELEMENTWISE_KERNEL(EqualsScalarForward) +REGISTER_MLU_ELEMENTWISE_KERNEL(AddForward) +REGISTER_MLU_ELEMENTWISE_KERNEL(AddBackward) +REGISTER_MLU_ELEMENTWISE_KERNEL(AddScalarForward) +REGISTER_MLU_ELEMENTWISE_KERNEL(AddScalarBackward) +REGISTER_MLU_ELEMENTWISE_KERNEL(SubForward) +REGISTER_MLU_ELEMENTWISE_KERNEL(SubBackward) +REGISTER_MLU_ELEMENTWISE_KERNEL(MulForward) +REGISTER_MLU_ELEMENTWISE_KERNEL(MulBackward) +REGISTER_MLU_ELEMENTWISE_KERNEL(MulScalarForward) +REGISTER_MLU_ELEMENTWISE_KERNEL(MulScalarBackward) +REGISTER_MLU_ELEMENTWISE_KERNEL(DivForward) +REGISTER_MLU_ELEMENTWISE_KERNEL(DivBackward) + +#undef REGISTER_MLU_ELEMENTWISE_KERNEL diff --git a/backends/mlu/src/kernels/elementwise.mlu b/backends/mlu/src/kernels/elementwise.mlu new file mode 100644 index 0000000..10da081 --- /dev/null +++ b/backends/mlu/src/kernels/elementwise.mlu @@ -0,0 +1,481 @@ +// 纯 BANG C kernel + 启动器（cncc 编译，只允许 C++11 级别语法，不含框架头）。 +// 框架集成见同目录 elementwise.cc。 +// 语义对应上游 src/kernels/cuda/elementwise.cu；约定： +// - 16 位浮点一律经 float 中转计算（FP32 中间精度）； +// - 广播反向用确定性 gather 累加（不用 atomic），FP32 累加。 + +#include +#include + +#include + +#include "kernels/mlu_kernel_api.h" + +namespace { +using infini_train::kernels::mlu::MluBinaryOp; +using infini_train::kernels::mlu::MluBroadcastMeta; +using infini_train::kernels::mlu::MluScalarType; +using infini_train::kernels::mlu::MluUnaryOp; + +constexpr uint32_t kDefaultTaskCount = 4; + +// 注意（实测）：本工具链下 taskIdX/coreId/taskDimX/coreDim 与预期语义不符，配合 +// cnrtFuncTypeUnion1 时循环写 GDRAM 只有首尾元素落盘。厂商示例（samples/BANG/.../ +// vectoradd/0_scalarAdd_single_core）用的是纯串行循环 + cnrtFuncTypeBlock。 +// 因此所有 kernel 统一改为单任务串行；性能优化属任务三，届时再按官方多核范式改写。 + +// 与 CUDA 版 CalcOffset 相同：把 out 的线性下标映射到广播输入的偏移 +__mlu_func__ 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) { + const int64_t out_index = (idx / out_strides[i]) % shape[i]; + offset += (shape[i] == 1 ? 0 : out_index) * strides[i]; + } + return offset; +} + +// ----------------------------- 一元 ----------------------------- + +__mlu_func__ float UnaryFwdEval(MluUnaryOp op, float x) { + switch (op) { + case MluUnaryOp::kNeg: + return -x; + case MluUnaryOp::kReciprocal: + return 1.0f / x; + case MluUnaryOp::kTanh: + return tanhf(x); + case MluUnaryOp::kRsqrt: + return 1.0f / sqrtf(x); + case MluUnaryOp::kExp: + return expf(x); + case MluUnaryOp::kLog: + return logf(x); + case MluUnaryOp::kSin: + return sinf(x); + case MluUnaryOp::kCos: + return cosf(x); + case MluUnaryOp::kSigmoid: + return 1.0f / (1.0f + expf(-x)); + } + return 0.0f; +} + +// aux 语义：tanh/exp/sigmoid 传 forward 的 output，其余传 input +__mlu_func__ float UnaryBwdEval(MluUnaryOp op, float g, float aux) { + switch (op) { + case MluUnaryOp::kNeg: + return -g; + case MluUnaryOp::kReciprocal: + return -g / (aux * aux); + case MluUnaryOp::kTanh: + return g * (1.0f - aux * aux); + case MluUnaryOp::kRsqrt: + return -0.5f * g / (aux * sqrtf(aux)); + case MluUnaryOp::kExp: + return g * aux; + case MluUnaryOp::kLog: + return g / aux; + case MluUnaryOp::kSin: + return g * cosf(aux); + case MluUnaryOp::kCos: + return -g * sinf(aux); + case MluUnaryOp::kSigmoid: + return g * aux * (1.0f - aux); + } + return 0.0f; +} + +template __mlu_global__ void UnaryForwardKernel(T *out, const T *in, MluUnaryOp op, int64_t n) { + for (int64_t i = 0; i < n; ++i) { out[i] = static_cast(UnaryFwdEval(op, static_cast(in[i]))); } +} + +template +__mlu_global__ void UnaryBackwardKernel(T *grad_in, const T *grad_out, const T *aux, MluUnaryOp op, int64_t n) { + for (int64_t i = 0; i < n; ++i) { + grad_in[i] = static_cast(UnaryBwdEval(op, static_cast(grad_out[i]), static_cast(aux[i]))); + } +} + +// 说明：设备侧 powf 在「单任务串行 + cnrtFuncTypeBlock」启动下已实测可用 +// （tests/math_selftest：powf(x,3) 与 host 参考最大绝对误差 1.91e-06，bad=0/1024）。 +// 历史上“powf 返回 NaN”的结论源于错误的多任务/Union1 启动配置，已作废。 +template +__mlu_global__ void PowForwardKernel(T *out, const T *in, float exponent, bool scalar_is_base, int64_t n) { + for (int64_t i = 0; i < n; ++i) { + const float x = static_cast(in[i]); + out[i] = static_cast(scalar_is_base ? powf(exponent, x) : powf(x, exponent)); + } +} + +template +__mlu_global__ void PowBackwardKernel(T *grad_in, const T *grad_out, const T *in, float exponent, bool scalar_is_base, + int64_t n) { + for (int64_t i = 0; i < n; ++i) { + const float g = static_cast(grad_out[i]); + const float x = static_cast(in[i]); + // scalar_is_base: d(base^x)/dx = base^x * ln(base) + // 否则: d(x^e)/dx = e * x^(e-1) + const float v + = scalar_is_base ? g * powf(exponent, x) * logf(exponent) : g * exponent * powf(x, exponent - 1.0f); + grad_in[i] = static_cast(v); + } +} + +// ----------------------------- 二元（广播） ----------------------------- + +__mlu_func__ float BinaryFwdEval(MluBinaryOp op, float a, float b) { + switch (op) { + case MluBinaryOp::kAdd: + return a + b; + case MluBinaryOp::kSub: + return a - b; + case MluBinaryOp::kMul: + return a * b; + case MluBinaryOp::kDiv: + return a / b; + } + return 0.0f; +} + +// side A 的 grad 公式 +__mlu_func__ float BinaryBwdEvalA(MluBinaryOp op, float g, float a, float b) { + switch (op) { + case MluBinaryOp::kAdd: + return g; + case MluBinaryOp::kSub: + return g; + case MluBinaryOp::kMul: + return g * b; + case MluBinaryOp::kDiv: + return g / b; + } + return 0.0f; +} + +// side B 的 grad 公式 +__mlu_func__ float BinaryBwdEvalB(MluBinaryOp op, float g, float a, float b) { + switch (op) { + case MluBinaryOp::kAdd: + return g; + case MluBinaryOp::kSub: + return -g; + case MluBinaryOp::kMul: + return g * a; + case MluBinaryOp::kDiv: + return -g * a / (b * b); + } + return 0.0f; +} + +template +__mlu_global__ void BinaryForwardKernel(T *out, const T *a, const T *b, MluBinaryOp op, MluBroadcastMeta meta, + int64_t n) { + for (int64_t i = 0; i < n; ++i) { + const int64_t a_off = CalcOffset(i, meta.ndim, meta.a_strides, meta.a_shape, meta.out_strides); + const int64_t b_off = CalcOffset(i, meta.ndim, meta.b_strides, meta.b_shape, meta.out_strides); + out[i] = static_cast(BinaryFwdEval(op, static_cast(a[a_off]), static_cast(b[b_off]))); + } +} + +// 无广播快路径：三侧形状一致，一一对应 +template +__mlu_global__ void BinaryBackwardNoBroadcastKernel(T *grad_a, T *grad_b, const T *grad_out, const T *a, const T *b, + MluBinaryOp op, int64_t n) { + for (int64_t i = 0; i < n; ++i) { + const float g = static_cast(grad_out[i]); + const float av = static_cast(a[i]); + const float bv = static_cast(b[i]); + grad_a[i] = static_cast(BinaryBwdEvalA(op, g, av, bv)); + grad_b[i] = static_cast(BinaryBwdEvalB(op, g, av, bv)); + } +} + +// 广播反向（gather 归约）：对 dst（= a 或 b 的 grad）的每个元素，枚举其在 out +// 中对应的 collapsed 下标组合并累加。kSideA=true 时 dst 对应 a，否则对应 b。 +template +__mlu_global__ void BinaryBackwardReduceKernel(T *grad_dst, const T *grad_out, const T *a, const T *b, MluBinaryOp op, + MluBroadcastMeta meta, int64_t numel_dst, int64_t numel_out) { + const int64_t *dst_shape = kSideA ? meta.a_shape : meta.b_shape; + const int64_t *dst_strides = kSideA ? meta.a_strides : meta.b_strides; + + for (int64_t i = 0; i < numel_dst; ++i) { + // 分解 i 到 dst（右对齐后）各维坐标 + int64_t coord[MLU_MAX_BROADCAST_DIMS]; + for (int d = 0; d < meta.ndim; ++d) { coord[d] = dst_shape[d] > 1 ? (i / dst_strides[d]) % dst_shape[d] : 0; } + // collapsed 空间 = dst_shape[d]==1 且 out_shape[d]>1 的维 + int64_t collapsed = 1; + for (int d = 0; d < meta.ndim; ++d) { + const int64_t out_extent = meta.a_shape[d] > meta.b_shape[d] ? meta.a_shape[d] : meta.b_shape[d]; + if (dst_shape[d] == 1 && out_extent > 1) { + collapsed *= out_extent; + } + } + float acc = 0.0f; + for (int64_t j = 0; j < collapsed; ++j) { + // 把 j 按各 collapsed 维的 extent 分解，得到完整 out 坐标 + int64_t rem = j; + int64_t out_coord[MLU_MAX_BROADCAST_DIMS]; + for (int d = 0; d < meta.ndim; ++d) { out_coord[d] = coord[d]; } + for (int d = meta.ndim - 1; d >= 0; --d) { + const int64_t out_extent = meta.a_shape[d] > meta.b_shape[d] ? meta.a_shape[d] : meta.b_shape[d]; + if (dst_shape[d] == 1 && out_extent > 1) { + out_coord[d] = rem % out_extent; + rem /= out_extent; + } + } + int64_t out_off = 0; + for (int d = 0; d < meta.ndim; ++d) { out_off += out_coord[d] * meta.out_strides[d]; } + const int64_t a_off + = kSideA ? i : CalcOffset(out_off, meta.ndim, meta.a_strides, meta.a_shape, meta.out_strides); + const int64_t b_off + = kSideA ? CalcOffset(out_off, meta.ndim, meta.b_strides, meta.b_shape, meta.out_strides) : i; + const float g = static_cast(grad_out[out_off]); + const float av = static_cast(a[a_off]); + const float bv = static_cast(b[b_off]); + acc += kSideA ? BinaryBwdEvalA(op, g, av, bv) : BinaryBwdEvalB(op, g, av, bv); + } + grad_dst[i] = static_cast(acc); + } + (void)numel_out; +} + +// ----------------------------- 标量二元 ----------------------------- + +template +__mlu_global__ void ScalarForwardKernel(T *out, const T *in, float scalar, MluBinaryOp op, int64_t n) { + for (int64_t i = 0; i < n; ++i) { out[i] = static_cast(BinaryFwdEval(op, static_cast(in[i]), scalar)); } +} + +template +__mlu_global__ void ScalarBackwardKernel(T *grad_in, const T *grad_out, float scalar, MluBinaryOp op, int64_t n) { + for (int64_t i = 0; i < n; ++i) { + const float g = static_cast(grad_out[i]); + // Add: grad = g；Mul: grad = g * scalar + grad_in[i] = static_cast(op == MluBinaryOp::kMul ? g * scalar : g); + } +} + +// ----------------------------- 相等比较 ----------------------------- + +template +__mlu_global__ void EqualsForwardKernel(bool *out, const T *a, const T *b, MluBroadcastMeta meta, int64_t n) { + for (int64_t i = 0; i < n; ++i) { + const int64_t a_off = CalcOffset(i, meta.ndim, meta.a_strides, meta.a_shape, meta.out_strides); + const int64_t b_off = CalcOffset(i, meta.ndim, meta.b_strides, meta.b_shape, meta.out_strides); + out[i] = a[a_off] == b[b_off]; + } +} + +template __mlu_global__ void EqualsScalarForwardKernel(bool *out, const T *in, T scalar, int64_t n) { + for (int64_t i = 0; i < n; ++i) { out[i] = in[i] == scalar; } +} + +// ----------------------------- 启动器 ----------------------------- + +const cnrtDim3_t kDim = {1, 1, 1}; + +template void UnaryForwardTyped(void *out, const void *in, MluUnaryOp op, int64_t n, cnrtQueue_t queue) { + UnaryForwardKernel<<>>(static_cast(out), static_cast(in), op, n); +} + +template +void UnaryBackwardTyped(void *grad_in, const void *grad_out, const void *aux, MluUnaryOp op, int64_t n, + cnrtQueue_t queue) { + UnaryBackwardKernel<<>>( + static_cast(grad_in), static_cast(grad_out), static_cast(aux), op, n); +} + +template +void PowForwardTyped(void *out, const void *in, float exponent, bool scalar_is_base, int64_t n, cnrtQueue_t queue) { + PowForwardKernel<<>>(static_cast(out), static_cast(in), exponent, + scalar_is_base, n); +} + +template +void PowBackwardTyped(void *grad_in, const void *grad_out, const void *in, float exponent, bool scalar_is_base, + int64_t n, cnrtQueue_t queue) { + PowBackwardKernel<<>>(static_cast(grad_in), + static_cast(grad_out), + static_cast(in), exponent, scalar_is_base, n); +} + +template +void BinaryForwardTyped(void *out, const void *a, const void *b, MluBinaryOp op, MluBroadcastMeta meta, int64_t n, + cnrtQueue_t queue) { + BinaryForwardKernel<<>>(static_cast(out), static_cast(a), + static_cast(b), op, meta, n); +} + +template +void BinaryBackwardTyped(void *grad_a, void *grad_b, const void *grad_out, const void *a, const void *b, MluBinaryOp op, + MluBroadcastMeta meta, int64_t numel_a, int64_t numel_b, int64_t numel_out, + cnrtQueue_t queue) { + bool same_shape = true; + for (int d = 0; d < meta.ndim; ++d) { + if (meta.a_shape[d] != meta.b_shape[d]) { + same_shape = false; + break; + } + } + if (same_shape) { + BinaryBackwardNoBroadcastKernel<<>>( + static_cast(grad_a), static_cast(grad_b), static_cast(grad_out), + static_cast(a), static_cast(b), op, numel_out); + return; + } + BinaryBackwardReduceKernel<<>>( + static_cast(grad_a), static_cast(grad_out), static_cast(a), + static_cast(b), op, meta, numel_a, numel_out); + BinaryBackwardReduceKernel<<>>( + static_cast(grad_b), static_cast(grad_out), static_cast(a), + static_cast(b), op, meta, numel_b, numel_out); +} + +template +void ScalarForwardTyped(void *out, const void *in, float scalar, MluBinaryOp op, int64_t n, cnrtQueue_t queue) { + ScalarForwardKernel + <<>>(static_cast(out), static_cast(in), scalar, op, n); +} + +template +void ScalarBackwardTyped(void *grad_in, const void *grad_out, float scalar, MluBinaryOp op, int64_t n, + cnrtQueue_t queue) { + ScalarBackwardKernel<<>>(static_cast(grad_in), + static_cast(grad_out), scalar, op, n); +} + +template +void EqualsForwardTyped(void *out, const void *a, const void *b, MluBroadcastMeta meta, int64_t n, cnrtQueue_t queue) { + EqualsForwardKernel<<>>(static_cast(out), static_cast(a), + static_cast(b), meta, n); +} + +template T ScalarFromDouble(double v) { + // 16 位浮点没有 double 构造函数（static_cast 歧义），经 float 中转 + if constexpr (std::is_same_v || std::is_same_v) { + return static_cast(static_cast(v)); + } else { + return static_cast(v); + } +} + +template +void EqualsScalarForwardTyped(void *out, const void *in, double scalar, int64_t n, cnrtQueue_t queue) { + // host 侧直接从 double 转换，整型比较不经 float 避免大数精度损失 + EqualsScalarForwardKernel<<>>( + static_cast(out), static_cast(in), ScalarFromDouble(scalar), n); +} + +// 浮点三型分发（FP64 已由 host 侧排除） +#define MLU_EW_FLOAT_SWITCH(fn, ...) \ + switch (type) { \ + case MluScalarType::kF16: \ + fn(__VA_ARGS__); \ + return; \ + case MluScalarType::kBF16: \ + fn(__VA_ARGS__); \ + return; \ + case MluScalarType::kF32: \ + fn(__VA_ARGS__); \ + return; \ + default: \ + return; \ + } + +// 全数值型分发（比较用；FP64 已由 host 侧排除） +#define MLU_EW_NUMERIC_SWITCH(fn, ...) \ + switch (type) { \ + case MluScalarType::kU8: \ + fn(__VA_ARGS__); \ + return; \ + case MluScalarType::kI8: \ + fn(__VA_ARGS__); \ + return; \ + case MluScalarType::kU16: \ + fn(__VA_ARGS__); \ + return; \ + case MluScalarType::kI16: \ + fn(__VA_ARGS__); \ + return; \ + case MluScalarType::kU32: \ + fn(__VA_ARGS__); \ + return; \ + case MluScalarType::kI32: \ + fn(__VA_ARGS__); \ + return; \ + case MluScalarType::kU64: \ + fn(__VA_ARGS__); \ + return; \ + case MluScalarType::kI64: \ + fn(__VA_ARGS__); \ + return; \ + case MluScalarType::kF16: \ + fn(__VA_ARGS__); \ + return; \ + case MluScalarType::kBF16: \ + fn(__VA_ARGS__); \ + return; \ + case MluScalarType::kF32: \ + fn(__VA_ARGS__); \ + return; \ + default: \ + return; \ + } +} // namespace + +namespace infini_train::kernels::mlu { + +void UnaryForwardLaunch(void *out, const void *in, MluUnaryOp op, MluScalarType type, int64_t num_elements, + cnrtQueue_t queue) { + MLU_EW_FLOAT_SWITCH(UnaryForwardTyped, out, in, op, num_elements, queue) +} + +void UnaryBackwardLaunch(void *grad_in, const void *grad_out, const void *aux, MluUnaryOp op, MluScalarType type, + int64_t num_elements, cnrtQueue_t queue) { + MLU_EW_FLOAT_SWITCH(UnaryBackwardTyped, grad_in, grad_out, aux, op, num_elements, queue) +} + +void PowForwardLaunch(void *out, const void *in, float exponent, bool scalar_is_base, MluScalarType type, + int64_t num_elements, cnrtQueue_t queue) { + MLU_EW_FLOAT_SWITCH(PowForwardTyped, out, in, exponent, scalar_is_base, num_elements, queue) +} + +void PowBackwardLaunch(void *grad_in, const void *grad_out, const void *in, float exponent, bool scalar_is_base, + MluScalarType type, int64_t num_elements, cnrtQueue_t queue) { + MLU_EW_FLOAT_SWITCH(PowBackwardTyped, grad_in, grad_out, in, exponent, scalar_is_base, num_elements, queue) +} + +void BinaryForwardLaunch(void *out, const void *a, const void *b, MluBinaryOp op, MluScalarType type, + MluBroadcastMeta meta, int64_t num_elements, cnrtQueue_t queue) { + MLU_EW_FLOAT_SWITCH(BinaryForwardTyped, out, a, b, op, meta, num_elements, queue) +} + +void BinaryBackwardLaunch(void *grad_a, void *grad_b, const void *grad_out, const void *a, const void *b, + MluBinaryOp op, MluScalarType type, MluBroadcastMeta meta, int64_t numel_a, int64_t numel_b, + int64_t numel_out, cnrtQueue_t queue) { + MLU_EW_FLOAT_SWITCH(BinaryBackwardTyped, grad_a, grad_b, grad_out, a, b, op, meta, numel_a, numel_b, numel_out, + queue) +} + +void ScalarForwardLaunch(void *out, const void *in, float scalar, MluBinaryOp op, MluScalarType type, + int64_t num_elements, cnrtQueue_t queue) { + MLU_EW_FLOAT_SWITCH(ScalarForwardTyped, out, in, scalar, op, num_elements, queue) +} + +void ScalarBackwardLaunch(void *grad_in, const void *grad_out, float scalar, MluBinaryOp op, MluScalarType type, + int64_t num_elements, cnrtQueue_t queue) { + MLU_EW_FLOAT_SWITCH(ScalarBackwardTyped, grad_in, grad_out, scalar, op, num_elements, queue) +} + +void EqualsForwardLaunch(void *out, const void *a, const void *b, MluScalarType type, MluBroadcastMeta meta, + int64_t num_elements, cnrtQueue_t queue) { + MLU_EW_NUMERIC_SWITCH(EqualsForwardTyped, out, a, b, meta, num_elements, queue) +} + +void EqualsScalarForwardLaunch(void *out, const void *in, double scalar, MluScalarType type, int64_t num_elements, + cnrtQueue_t queue) { + MLU_EW_NUMERIC_SWITCH(EqualsScalarForwardTyped, out, in, scalar, num_elements, queue) +} + +} // namespace infini_train::kernels::mlu diff --git a/backends/mlu/src/kernels/embedding.cc b/backends/mlu/src/kernels/embedding.cc new file mode 100644 index 0000000..c9737f4 --- /dev/null +++ b/backends/mlu/src/kernels/embedding.cc @@ -0,0 +1,90 @@ +// 对应上游 src/kernels/cuda/embedding.cu 的 host 集成部分；设备 kernel 在 embedding.mlu。 + +#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/device.h" +#include "infini_train/include/dispatcher.h" +#include "infini_train/include/tensor.h" + +#include "kernels/mlu_kernel_api.h" +#include "runtime/mlu_debug.h" +#include "runtime/mlu_dispatch.h" +#include "runtime/mlu_runtime_common.h" + +namespace infini_train::kernels::mlu { +namespace { +cnrtQueue_t CurrentQueue(const Device &device) { + return dynamic_cast( + infini_train::core::GetDeviceGuardImpl(device.type())->GetStream(device)) + ->mlu_queue(); +} +} // namespace + +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); + + const int64_t batch_size = input->Dims().size() == 2 ? input->Dims()[0] : 1; + const int64_t max_seqlen = input->Dims().size() == 2 ? input->Dims()[1] : input->Dims()[0]; + const int64_t vocab_size = weight->Dims()[0]; + const int64_t embed_dim = weight->Dims()[1]; + auto output_dims = input->Dims(); + output_dims.push_back(embed_dim); + + auto output = std::make_shared(output_dims, weight->Dtype(), input->GetDevice()); + const auto queue = CurrentQueue(input->GetDevice()); + + core::mlu::DispatchMluFunc( + weight->Dtype(), + [=]() { + EmbeddingForwardLaunch(input->DataPtr(), output->DataPtr(), weight->DataPtr(), batch_size * max_seqlen, + embed_dim, vocab_size, MluTypeTag(), queue); + ::infini_train::core::mlu::DebugSyncQueue(queue, "EmbeddingForward"); + ::infini_train::core::mlu::DebugCheck(output, "EmbeddingForward"); + }, + "MLU EmbeddingForward"); + return output; +} + +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); + + const int64_t vocab_size = weight_dims[0]; + const int64_t embedding_dim = weight_dims[1]; + CHECK_EQ(input->Dims().size() + 1, grad_output->Dims().size()); + for (size_t 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 grad_weight = std::make_shared(weight_dims, grad_output->Dtype(), grad_output->GetDevice()); + // atomic 累加前必须清零 + grad_weight->Fill(0.0); + const auto queue = CurrentQueue(input->GetDevice()); + + core::mlu::DispatchMluFunc( + grad_output->Dtype(), + [=]() { + EmbeddingBackwardLaunch(input->DataPtr(), grad_output->DataPtr(), grad_weight->DataPtr(), + input->NumElements(), embedding_dim, vocab_size, MluTypeTag(), queue); + ::infini_train::core::mlu::DebugSyncQueue(queue, "EmbeddingBackward"); + }, + "MLU EmbeddingBackward"); + return grad_weight; +} +} // namespace infini_train::kernels::mlu + +#define REGISTER_MLU_EMBEDDING_KERNEL(kernel_name) \ + REGISTER_KERNEL(infini_train::Device::DeviceType::kPrivateUse1, kernel_name, \ + infini_train::kernels::mlu::kernel_name) + +REGISTER_MLU_EMBEDDING_KERNEL(EmbeddingForward) +REGISTER_MLU_EMBEDDING_KERNEL(EmbeddingBackward) + +#undef REGISTER_MLU_EMBEDDING_KERNEL diff --git a/backends/mlu/src/kernels/embedding.mlu b/backends/mlu/src/kernels/embedding.mlu new file mode 100644 index 0000000..8743e3e --- /dev/null +++ b/backends/mlu/src/kernels/embedding.mlu @@ -0,0 +1,108 @@ +// 纯 BANG C kernel + 启动器（cncc 编译，只允许 C++11 级别语法，不含框架头）。 +// 框架集成见同目录 embedding.cc。 +// 语义对应上游 src/kernels/cuda/embedding.cu。 + +#include + +#include + +#include "kernels/mlu_kernel_api.h" + +namespace { +using infini_train::kernels::mlu::MluScalarType; + +constexpr uint32_t kDefaultTaskCount = 4; + +// 注意（实测）：本工具链下 taskIdX/coreId/taskDimX/coreDim 与预期语义不符，配合 +// cnrtFuncTypeUnion1 时循环写 GDRAM 只有首尾元素落盘。厂商示例（samples/BANG/.../ +// vectoradd/0_scalarAdd_single_core）用的是纯串行循环 + cnrtFuncTypeBlock。 +// 因此所有 kernel 统一改为单任务串行；性能优化属任务三，届时再按官方多核范式改写。 + +template +__mlu_global__ void EmbeddingForwardKernel(const int64_t *input, T *output, const T *weight, int64_t num_tokens, + int64_t embed_dim, int64_t vocab_size) { + const int64_t total = num_tokens * embed_dim; + + for (int64_t idx = 0; idx < total; ++idx) { + const int64_t t = idx / embed_dim; + const int64_t c = idx % embed_dim; + const int64_t ix = input[t]; + if (ix < 0 || ix >= vocab_size) { + continue; + } + output[t * embed_dim + c] = weight[ix * embed_dim + c]; + } +} + +template +__mlu_global__ void EmbeddingBackwardKernel(const int64_t *input, const T *grad_output, T *grad_weight, + int64_t num_tokens, int64_t embed_dim, int64_t vocab_size) { + // 不用 __bang_atomic_add：其参数语义（dst/src1/src2 指针含义、是否要求 NRAM 侧 + // 指针）在 MLU590 上实测会触发队列错误（CN_INVOKE_ERROR）。 + // 改为"一个核负责若干词表行，串行扫所有 token"：每行只有一个写者，无需原子； + // 累加顺序按 token 升序，结果确定、可复现（精度优先，性能属任务三）。 + // TODO(BF16-MUST-FIX, 001-P2 F07): 当前内层把 float 结果立即写回 T，FP32 下无损， + // 但 BF16/FP16 会逐 token 舍入，不能称为“FP32 中间累加”。BF16 阶段必须改为 + // 块内 float 累加、最后一次转换。 + for (int64_t row = 0; row < vocab_size; ++row) { + T *dst = grad_weight + row * embed_dim; + for (int64_t t = 0; t < num_tokens; ++t) { + if (input[t] != row) { + continue; + } + const T *src = grad_output + t * embed_dim; + for (int64_t j = 0; j < embed_dim; ++j) { + dst[j] = static_cast(static_cast(dst[j]) + static_cast(src[j])); + } + } + } +} + +template +void EmbeddingForwardTyped(const void *input, void *output, const void *weight, int64_t num_tokens, int64_t embed_dim, + int64_t vocab_size, cnrtQueue_t queue) { + const cnrtDim3_t dim = {1, 1, 1}; + EmbeddingForwardKernel + <<>>(static_cast(input), static_cast(output), + static_cast(weight), num_tokens, embed_dim, vocab_size); +} + +template +void EmbeddingBackwardTyped(const void *input, const void *grad_output, void *grad_weight, int64_t num_tokens, + int64_t embed_dim, int64_t vocab_size, cnrtQueue_t queue) { + const cnrtDim3_t dim = {1, 1, 1}; + EmbeddingBackwardKernel + <<>>(static_cast(input), static_cast(grad_output), + static_cast(grad_weight), num_tokens, embed_dim, vocab_size); +} + +#define MLU_EMBED_FLOAT_SWITCH(fn, ...) \ + switch (type) { \ + case MluScalarType::kF16: \ + fn(__VA_ARGS__); \ + return; \ + case MluScalarType::kBF16: \ + fn(__VA_ARGS__); \ + return; \ + case MluScalarType::kF32: \ + fn(__VA_ARGS__); \ + return; \ + default: \ + return; \ + } +} // namespace + +namespace infini_train::kernels::mlu { + +void EmbeddingForwardLaunch(const void *input, void *output, const void *weight, int64_t num_tokens, int64_t embed_dim, + int64_t vocab_size, MluScalarType type, cnrtQueue_t queue) { + MLU_EMBED_FLOAT_SWITCH(EmbeddingForwardTyped, input, output, weight, num_tokens, embed_dim, vocab_size, queue) +} + +void EmbeddingBackwardLaunch(const void *input, const void *grad_output, void *grad_weight, int64_t num_tokens, + int64_t embed_dim, int64_t vocab_size, MluScalarType type, cnrtQueue_t queue) { + MLU_EMBED_FLOAT_SWITCH(EmbeddingBackwardTyped, input, grad_output, grad_weight, num_tokens, embed_dim, vocab_size, + queue) +} + +} // namespace infini_train::kernels::mlu diff --git a/backends/mlu/src/kernels/fill.cc b/backends/mlu/src/kernels/fill.cc new file mode 100644 index 0000000..317e00c --- /dev/null +++ b/backends/mlu/src/kernels/fill.cc @@ -0,0 +1,69 @@ +// 对应上游 src/kernels/cuda/fill.cu 的 host 集成部分；设备 kernel 在 fill.mlu。 + +#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/scalar.h" +#include "infini_train/include/tensor.h" + +#include "kernels/mlu_kernel_api.h" +#include "runtime/mlu_dispatch.h" +#include "runtime/mlu_runtime_common.h" + +namespace infini_train::kernels::mlu { + +namespace { + +// Scalar::to() 对 MLU 原生 16 位类型不能直接用：host 侧 __bang_bfloat16 +// 没有 double 构造函数（static_cast 在 float/int 等整型 ctor 间歧义）。 +// 按 scalar.h 中 TODO 的建议，后端自定义 scalar_cast，半精度统一走 +// double -> float -> half/bf16 两级转换，与框架 to() 的舍入路径一致。 +template T MluScalarCast(const Scalar &scalar) { + if constexpr (std::is_same_v || std::is_same_v) { + return T(scalar.to()); + } else { + return scalar.to(); + } +} + +} // namespace + +void Fill(std::shared_ptr tensor, Scalar scalar) { + auto device = tensor->GetDevice(); + const auto queue = dynamic_cast( + infini_train::core::GetDeviceGuardImpl(device.type())->GetStream(device)) + ->mlu_queue(); + + core::mlu::DispatchMluFunc( + tensor->Dtype(), + [=]() { + const T value = MluScalarCast(scalar); + static_assert(sizeof(T) <= sizeof(uint64_t)); + uint64_t value_bits = 0; + std::memcpy(&value_bits, &value, sizeof(T)); + // 零填充快路径：单核标量 kernel 写 154MB（如 embedding 梯度表）要跑分钟级， + // 实测会触发设备错误把队列搞崩；交零直接用 cnrtMemset（同步） + if (value_bits == 0) { + CNRT_CHECK(cnrtMemset(tensor->DataPtr(), 0, static_cast(tensor->NumElements()) * sizeof(T))); + return; + } + FillLaunch(tensor->DataPtr(), value_bits, MluTypeTag(), static_cast(tensor->NumElements()), + queue); + ::infini_train::core::mlu::DebugSyncQueue(queue, "Fill"); + }, + "MLU Fill"); +} +} // namespace infini_train::kernels::mlu + +#define REGISTER_MLU_FILL_KERNEL(kernel_name) \ + REGISTER_KERNEL(infini_train::Device::DeviceType::kPrivateUse1, kernel_name, \ + infini_train::kernels::mlu::kernel_name) + +REGISTER_MLU_FILL_KERNEL(Fill) + +#undef REGISTER_MLU_FILL_KERNEL diff --git a/backends/mlu/src/kernels/fill.mlu b/backends/mlu/src/kernels/fill.mlu new file mode 100644 index 0000000..3d3b95f --- /dev/null +++ b/backends/mlu/src/kernels/fill.mlu @@ -0,0 +1,71 @@ +// 纯 BANG C kernel + 启动器（cncc 编译，只允许 C++11 级别语法，不含框架头）。 +// 框架集成见同目录 fill.cc。 + +#include +#include + +#include + +#include "kernels/mlu_kernel_api.h" + +namespace { +// TODO(perf): 任务数先取保守值（4 个 union1 任务 = 16 核），性能是任务三的事 +constexpr uint32_t kDefaultTaskCount = 4; + +template __mlu_global__ void FillKernel(T *data, T value, int64_t size) { + + for (int64_t idx = 0; idx < size; ++idx) { data[idx] = value; } +} + +template void FillTyped(void *data, uint64_t value_bits, int64_t num_elements, cnrtQueue_t queue) { + T value; + memcpy(&value, &value_bits, sizeof(T)); + const cnrtDim3_t dim = {1, 1, 1}; + FillKernel<<>>(static_cast(data), value, num_elements); +} +} // namespace + +namespace infini_train::kernels::mlu { + +void FillLaunch(void *data, uint64_t value_bits, MluScalarType type, int64_t num_elements, cnrtQueue_t queue) { + switch (type) { + case MluScalarType::kU8: + FillTyped(data, value_bits, num_elements, queue); + break; + case MluScalarType::kI8: + FillTyped(data, value_bits, num_elements, queue); + break; + case MluScalarType::kU16: + FillTyped(data, value_bits, num_elements, queue); + break; + case MluScalarType::kI16: + FillTyped(data, value_bits, num_elements, queue); + break; + case MluScalarType::kU32: + FillTyped(data, value_bits, num_elements, queue); + break; + case MluScalarType::kI32: + FillTyped(data, value_bits, num_elements, queue); + break; + case MluScalarType::kU64: + FillTyped(data, value_bits, num_elements, queue); + break; + case MluScalarType::kI64: + FillTyped(data, value_bits, num_elements, queue); + break; + case MluScalarType::kF16: + FillTyped(data, value_bits, num_elements, queue); + break; + case MluScalarType::kBF16: + FillTyped(data, value_bits, num_elements, queue); + break; + case MluScalarType::kF32: + FillTyped(data, value_bits, num_elements, queue); + break; + default: + // MLU 设备不支持 FP64；bool 与 CUDA 后端一致不 dispatch（见 fill.cc） + break; + } +} + +} // namespace infini_train::kernels::mlu diff --git a/backends/mlu/src/kernels/layernorm.cc b/backends/mlu/src/kernels/layernorm.cc new file mode 100644 index 0000000..4408266 --- /dev/null +++ b/backends/mlu/src/kernels/layernorm.cc @@ -0,0 +1,104 @@ +// 对应上游 src/kernels/cuda/layernorm.cu 的 host 集成部分；设备 kernel 在 layernorm.mlu。 + +#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/device.h" +#include "infini_train/include/dispatcher.h" +#include "infini_train/include/tensor.h" + +#include "kernels/mlu_kernel_api.h" +#include "runtime/mlu_debug.h" +#include "runtime/mlu_dispatch.h" +#include "runtime/mlu_runtime_common.h" + +namespace infini_train::kernels::mlu { +namespace { +cnrtQueue_t CurrentQueue(const Device &device) { + return dynamic_cast( + infini_train::core::GetDeviceGuardImpl(device.type())->GetStream(device)) + ->mlu_queue(); +} +} // namespace + +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 int64_t batch_size = input->Dims()[0]; + const int64_t max_seqlen = input->Dims()[1]; + const int64_t 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()); + + const auto queue = CurrentQueue(input->GetDevice()); + core::mlu::DispatchMluFunc( + dtype, + [=]() { + LayerNormForwardLaunch(input->DataPtr(), weight->DataPtr(), bias->DataPtr(), + static_cast(mean->DataPtr()), static_cast(rstd->DataPtr()), + output->DataPtr(), eps, batch_size * max_seqlen, embed_dim, MluTypeTag(), queue); + ::infini_train::core::mlu::DebugSyncQueue(queue, "LayerNormForward"); + ::infini_train::core::mlu::DebugCheck(output, "LayerNormForward"); + }, + "MLU LayerNormForward"); + return {output, mean, rstd}; +} + +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 int64_t batch_size = input->Dims()[0]; + const int64_t max_seqlen = input->Dims()[1]; + const int64_t embed_dim = input->Dims()[2]; + const int64_t num_tokens = 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()); + + const auto queue = CurrentQueue(grad_output->GetDevice()); + core::mlu::DispatchMluFunc( + dtype, + [=]() { + LayerNormBackwardLaunch( + input->DataPtr(), grad_output->DataPtr(), static_cast(mean->DataPtr()), + static_cast(rstd->DataPtr()), weight->DataPtr(), grad_input->DataPtr(), + grad_weight->DataPtr(), grad_bias->DataPtr(), num_tokens, embed_dim, MluTypeTag(), queue); + ::infini_train::core::mlu::DebugSyncQueue(queue, "LayerNormBackward"); + }, + "MLU LayerNormBackward"); + return {grad_input, grad_weight, grad_bias}; +} +} // namespace infini_train::kernels::mlu + +#define REGISTER_MLU_LAYERNORM_KERNEL(kernel_name) \ + REGISTER_KERNEL(infini_train::Device::DeviceType::kPrivateUse1, kernel_name, \ + infini_train::kernels::mlu::kernel_name) + +REGISTER_MLU_LAYERNORM_KERNEL(LayerNormForward) +REGISTER_MLU_LAYERNORM_KERNEL(LayerNormBackward) + +#undef REGISTER_MLU_LAYERNORM_KERNEL diff --git a/backends/mlu/src/kernels/layernorm.mlu b/backends/mlu/src/kernels/layernorm.mlu new file mode 100644 index 0000000..3905999 --- /dev/null +++ b/backends/mlu/src/kernels/layernorm.mlu @@ -0,0 +1,159 @@ +// 纯 BANG C kernel + 启动器（cncc 编译，只允许 C++11 级别语法，不含框架头）。 +// 框架集成见同目录 layernorm.cc。 +// 语义对应上游 src/kernels/cuda/layernorm.cu。 +// 简化策略（正确性优先，性能属任务三）：一行（token）分配一个核串行归约， +// FP32 累加、顺序确定；不引入跨核 SRAM block reduce。 + +#include + +#include + +#include "kernels/mlu_kernel_api.h" + +namespace { +using infini_train::kernels::mlu::MluScalarType; + +constexpr uint32_t kDefaultTaskCount = 4; + +// 注意（实测）：本工具链下 taskIdX/coreId/taskDimX/coreDim 与预期语义不符，配合 +// cnrtFuncTypeUnion1 时循环写 GDRAM 只有首尾元素落盘。厂商示例（samples/BANG/.../ +// vectoradd/0_scalarAdd_single_core）用的是纯串行循环 + cnrtFuncTypeBlock。 +// 因此所有 kernel 统一改为单任务串行；性能优化属任务三，届时再按官方多核范式改写。 + +template +__mlu_global__ void LayerNormForwardKernel(const T *input, const T *weight, const T *bias, float *mean_out, + float *rstd_out, T *output, float eps, int64_t num_tokens, + int64_t embed_dim) { + for (int64_t token = 0; token < num_tokens; ++token) { + const T *x = input + token * embed_dim; + T *y = output + token * embed_dim; + + float sum = 0.0f; + float sqsum = 0.0f; + for (int64_t i = 0; i < embed_dim; ++i) { + const float val = static_cast(x[i]); + sum += val; + sqsum += val * val; + } + const float mean = sum / embed_dim; + const float var = sqsum / embed_dim - mean * mean; + const float rstd = 1.0f / sqrtf(var + eps); + if (mean_out) { + mean_out[token] = mean; + } + if (rstd_out) { + rstd_out[token] = rstd; + } + + for (int64_t i = 0; i < embed_dim; ++i) { + const float norm = (static_cast(x[i]) - mean) * rstd; + y[i] = static_cast(norm * static_cast(weight[i]) + static_cast(bias[i])); + } + } +} + +template +__mlu_global__ void LayerNormInputGradKernel(const T *input, const T *grad_output, const float *mean, const float *rstd, + const T *weight, T *grad_input, int64_t num_tokens, int64_t embed_dim) { + for (int64_t token = 0; token < num_tokens; ++token) { + const T *x = input + token * embed_dim; + const T *g = grad_output + token * embed_dim; + T *gi = grad_input + token * embed_dim; + const float mean_val = mean[token]; + const float rstd_val = rstd[token]; + + float dnorm_mean = 0.0f; + float dnorm_norm_mean = 0.0f; + for (int64_t i = 0; i < embed_dim; ++i) { + const float dnorm = static_cast(weight[i]) * static_cast(g[i]); + const float norm = (static_cast(x[i]) - mean_val) * rstd_val; + dnorm_mean += dnorm; + dnorm_norm_mean += dnorm * norm; + } + dnorm_mean /= embed_dim; + dnorm_norm_mean /= embed_dim; + + for (int64_t i = 0; i < embed_dim; ++i) { + const float norm = (static_cast(x[i]) - mean_val) * rstd_val; + gi[i] = static_cast( + (static_cast(weight[i]) * static_cast(g[i]) - dnorm_mean - norm * dnorm_norm_mean) + * rstd_val); + } + } +} + +// 一个核负责一个 feature，串行扫所有 token（确定性顺序） +template +__mlu_global__ void LayerNormParameterGradKernel(const T *input, const T *grad_output, const float *mean, + const float *rstd, T *grad_weight, T *grad_bias, int64_t num_tokens, + int64_t embed_dim) { + for (int64_t feature = 0; feature < embed_dim; ++feature) { + float gw = 0.0f; + float gb = 0.0f; + for (int64_t token = 0; token < num_tokens; ++token) { + const int64_t offset = token * embed_dim + feature; + const float g = static_cast(grad_output[offset]); + const float norm = (static_cast(input[offset]) - mean[token]) * rstd[token]; + gw += g * norm; + gb += g; + } + grad_weight[feature] = static_cast(gw); + grad_bias[feature] = static_cast(gb); + } +} + +template +void LayerNormForwardTyped(const void *input, const void *weight, const void *bias, float *mean_out, float *rstd_out, + void *output, float eps, int64_t num_tokens, int64_t embed_dim, cnrtQueue_t queue) { + const cnrtDim3_t dim = {1, 1, 1}; + LayerNormForwardKernel<<>>( + static_cast(input), static_cast(weight), static_cast(bias), mean_out, rstd_out, + static_cast(output), eps, num_tokens, embed_dim); +} + +template +void LayerNormBackwardTyped(const void *input, const void *grad_output, const float *mean, const float *rstd, + const void *weight, void *grad_input, void *grad_weight, void *grad_bias, + int64_t num_tokens, int64_t embed_dim, cnrtQueue_t queue) { + const cnrtDim3_t dim = {1, 1, 1}; + LayerNormInputGradKernel<<>>( + static_cast(input), static_cast(grad_output), mean, rstd, static_cast(weight), + static_cast(grad_input), num_tokens, embed_dim); + LayerNormParameterGradKernel<<>>( + static_cast(input), static_cast(grad_output), mean, rstd, static_cast(grad_weight), + static_cast(grad_bias), num_tokens, embed_dim); +} + +#define MLU_LN_FLOAT_SWITCH(fn, ...) \ + switch (type) { \ + case MluScalarType::kF16: \ + fn(__VA_ARGS__); \ + return; \ + case MluScalarType::kBF16: \ + fn(__VA_ARGS__); \ + return; \ + case MluScalarType::kF32: \ + fn(__VA_ARGS__); \ + return; \ + default: \ + return; \ + } +} // namespace + +namespace infini_train::kernels::mlu { + +void LayerNormForwardLaunch(const void *input, const void *weight, const void *bias, float *mean_out, float *rstd_out, + void *output, float eps, int64_t num_tokens, int64_t embed_dim, MluScalarType type, + cnrtQueue_t queue) { + MLU_LN_FLOAT_SWITCH(LayerNormForwardTyped, input, weight, bias, mean_out, rstd_out, output, eps, num_tokens, + embed_dim, queue) +} + +void LayerNormBackwardLaunch(const void *input, const void *grad_output, const float *mean, const float *rstd, + const void *weight, void *grad_input, void *grad_weight, void *grad_bias, + int64_t num_tokens, int64_t embed_dim, MluScalarType type, cnrtQueue_t queue) { + MLU_LN_FLOAT_SWITCH(LayerNormBackwardTyped, input, grad_output, mean, rstd, weight, grad_input, grad_weight, + grad_bias, num_tokens, embed_dim, queue) +} + +} // namespace infini_train::kernels::mlu diff --git a/backends/mlu/src/kernels/linear.cc b/backends/mlu/src/kernels/linear.cc new file mode 100644 index 0000000..8e77bfc --- /dev/null +++ b/backends/mlu/src/kernels/linear.cc @@ -0,0 +1,318 @@ +// Linear / Matmul 的 MLU host 侧集成。 +// +// NOTE(MLU-FALLBACK) 设备 GEMM 尚未接入：本文件目前用 D2H -> host 分块 GEMM -> H2D +// 作为临时实现，只为先把 FP32 数值链路跑通；正式验收要求替换为设备端 GEMM +// （例如 CNNL，且 handle 必须绑定到当前 stream）。 +// 精度：host 侧为 float 顺序累加（k 递增），与分块累加的差异仅来自舍入。 + +#include +#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/device.h" +#include "infini_train/include/dispatcher.h" +#include "infini_train/include/tensor.h" + +#include "runtime/mlu_debug.h" +#include "runtime/mlu_dispatch.h" +#include "runtime/mlu_runtime_common.h" + +namespace infini_train::kernels::mlu { +namespace { + +// C(m,n) = op(A)(m,k) * op(B)(k,n) + beta * C(m,n) +// 行主序；A 的 lda 为行跨度，B 的 ldb 为行跨度。alpha 固定为 1（本项目用不到）。 +void HostGemm(bool trans_a, bool trans_b, int64_t m, int64_t n, int64_t k, const float *a, int64_t lda, const float *b, + int64_t ldb, float beta, float *c, int64_t ldc) { + // i-k-j 序：内层沿 n 连续访问 B、C，缓存友好；k 递增累加保证顺序确定性 +#pragma omp parallel for schedule(static) + for (int64_t i = 0; i < m; ++i) { + float *c_row = c + i * ldc; + for (int64_t j = 0; j < n; ++j) { c_row[j] = beta == 0.0f ? 0.0f : beta * c_row[j]; } + for (int64_t p = 0; p < k; ++p) { + const float av = trans_a ? a[p * lda + i] : a[i * lda + p]; + if (av == 0.0f) { + continue; + } + const float *b_row = trans_b ? (b + p) : (b + p * ldb); + const int64_t b_step = trans_b ? ldb : 1; + for (int64_t j = 0; j < n; ++j) { c_row[j] += av * b_row[j * b_step]; } + } + } +} + +// 设备张量 <-> host float 缓冲（16 位浮点经 float 中转） +std::vector ToHostFloat(const std::shared_ptr &t, Device::DeviceType type) { + const size_t num = t->NumElements(); + std::vector host(num); + auto *impl = infini_train::core::GetDeviceGuardImpl(type); + switch (t->Dtype()) { + case DataType::kFLOAT32: + impl->Memcpy(host.data(), t->DataPtr(), num * sizeof(float), infini_train::core::MemcpyKind::kD2H); + break; + case DataType::kFLOAT16: + case DataType::kBFLOAT16: { + std::vector raw(num); + impl->Memcpy(raw.data(), t->DataPtr(), num * sizeof(uint16_t), infini_train::core::MemcpyKind::kD2H); + for (size_t i = 0; i < num; ++i) { + if (t->Dtype() == DataType::kFLOAT16) { + host[i] = infini_train::detail::Fp16BitsToFloat(raw[i]); + } else { + host[i] = infini_train::detail::Bf16BitsToFloat(raw[i]); + } + } + break; + } + default: + LOG(FATAL) << "MLU GEMM fallback: unsupported dtype " << static_cast(t->Dtype()); + } + return host; +} + +void FromHostFloat(const std::vector &host, const std::shared_ptr &t, Device::DeviceType type) { + const size_t num = t->NumElements(); + auto *impl = infini_train::core::GetDeviceGuardImpl(type); + switch (t->Dtype()) { + case DataType::kFLOAT32: + impl->Memcpy(t->DataPtr(), host.data(), num * sizeof(float), infini_train::core::MemcpyKind::kH2D); + break; + case DataType::kFLOAT16: + case DataType::kBFLOAT16: { + std::vector raw(num); + for (size_t i = 0; i < num; ++i) { + raw[i] = t->Dtype() == DataType::kFLOAT16 ? infini_train::detail::FloatToFp16Bits(host[i]) + : infini_train::detail::FloatToBf16Bits(host[i]); + } + impl->Memcpy(t->DataPtr(), raw.data(), num * sizeof(uint16_t), infini_train::core::MemcpyKind::kH2D); + break; + } + default: + LOG(FATAL) << "MLU GEMM fallback: unsupported dtype " << static_cast(t->Dtype()); + } +} + +int64_t NumElementsOf(const std::vector &dims) { + return std::accumulate(dims.begin(), dims.end(), 1, std::multiplies()); +} + +} // namespace + +// ----------------------------------------------------------------------------- +// Linear +// ----------------------------------------------------------------------------- + +std::shared_ptr LinearForward(const std::shared_ptr &input, const std::shared_ptr &weight, + bool transpose, const std::shared_ptr &bias) { + const auto &input_dims = input->Dims(); + CHECK_GE(input_dims.size(), 2); + 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]); + const int64_t out_features = weight_dims[transpose ? 0 : 1]; + + if (bias) { + CHECK_EQ(bias->Dims().size(), 1); + CHECK_EQ(bias->Dims()[0], out_features); + } + + auto output_dims = input_dims; + *output_dims.rbegin() = out_features; + auto output = std::make_shared(output_dims, input->Dtype(), input->GetDevice()); + + const int64_t rows = NumElementsOf(input_dims) / in_features; + auto in_host = ToHostFloat(input, input->GetDevice().type()); + auto w_host = ToHostFloat(weight, weight->GetDevice().type()); + std::vector out_host(static_cast(rows * out_features)); + + // output = input(rows,in) * weight^T(in,out)（transpose=true 时 weight 为 (out,in)） + HostGemm(false, transpose, rows, out_features, in_features, in_host.data(), in_features, w_host.data(), + transpose ? in_features : out_features, 0.0f, out_host.data(), out_features); + + if (bias) { + auto b_host = ToHostFloat(bias, bias->GetDevice().type()); + for (int64_t r = 0; r < rows; ++r) { + for (int64_t j = 0; j < out_features; ++j) { out_host[r * out_features + j] += b_host[j]; } + } + } + + FromHostFloat(out_host, output, output->GetDevice().type()); + if (output->Dtype() == DataType::kFLOAT32) { + ::infini_train::core::mlu::DebugCheckTensor(output->DataPtr(), output->NumElements(), + static_cast(output->Dtype()), "LinearForward"); + } + return output; +} + +std::shared_ptr LinearBackwardInput(const std::shared_ptr &weight, + const std::shared_ptr &grad_output, bool transpose, + int64_t in_features, int64_t out_features, + const std::vector &input_dims) { + auto grad_input = std::make_shared(input_dims, grad_output->Dtype(), grad_output->GetDevice()); + const int64_t rows = NumElementsOf(input_dims) / in_features; + + auto g_host = ToHostFloat(grad_output, grad_output->GetDevice().type()); + auto w_host = ToHostFloat(weight, weight->GetDevice().type()); + std::vector gi_host(static_cast(rows * in_features)); + + if (transpose) { + // weight 为 [out, in]：grad_input(rows,in) = grad_output(rows,out) * weight(out,in) + HostGemm(false, false, rows, in_features, out_features, g_host.data(), out_features, w_host.data(), in_features, + 0.0f, gi_host.data(), in_features); + } else { + // weight 为 [in, out]：grad_input(rows,in) = grad_output(rows,out) * weight(in,out)^T + HostGemm(false, true, rows, in_features, out_features, g_host.data(), out_features, w_host.data(), out_features, + 0.0f, gi_host.data(), in_features); + } + + FromHostFloat(gi_host, grad_input, grad_input->GetDevice().type()); + 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 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, grad_output->Dtype(), grad_output->GetDevice()); + const int64_t rows = input->NumElements() / in_features; + + auto in_host = ToHostFloat(input, input->GetDevice().type()); + auto g_host = ToHostFloat(grad_output, grad_output->GetDevice().type()); + std::vector gw_host(static_cast(out_features * in_features)); + + if (transpose) { + // grad_weight(out,in) = grad_output^T(out,rows) * input(rows,in) + HostGemm(true, false, out_features, in_features, rows, g_host.data(), out_features, in_host.data(), in_features, + 0.0f, gw_host.data(), in_features); + } else { + // grad_weight(in,out) = input^T(in,rows) * grad_output(rows,out) + HostGemm(true, false, in_features, out_features, rows, in_host.data(), in_features, g_host.data(), out_features, + 0.0f, gw_host.data(), out_features); + } + + FromHostFloat(gw_host, grad_weight, grad_weight->GetDevice().type()); + return grad_weight; +} + +std::shared_ptr LinearBackwardBias(const std::shared_ptr &grad_output, int64_t out_features) { + auto grad_bias + = std::make_shared(std::vector{out_features}, grad_output->Dtype(), grad_output->GetDevice()); + const int64_t rows = grad_output->NumElements() / out_features; + + auto g_host = ToHostFloat(grad_output, grad_output->GetDevice().type()); + std::vector gb_host(static_cast(out_features), 0.0f); + for (int64_t r = 0; r < rows; ++r) { + for (int64_t j = 0; j < out_features; ++j) { gb_host[j] += g_host[r * out_features + j]; } + } + + FromHostFloat(gb_host, grad_bias, grad_bias->GetDevice().type()); + return grad_bias; +} + +// ----------------------------------------------------------------------------- +// Matmul（批量，最后两维为矩阵，batch 维必须一致） +// ----------------------------------------------------------------------------- + +std::shared_ptr MatmulForward(const std::shared_ptr &input, const std::shared_ptr &other) { + const auto &input_dims = input->Dims(); + const auto &other_dims = other->Dims(); + CHECK_GE(input_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 = NumElementsOf(std::vector(input_dims.begin(), input_dims.end() - 2)); + + auto output_dims = input_dims; + output_dims.back() = n; + auto output = std::make_shared(output_dims, input->Dtype(), input->GetDevice()); + + auto a_host = ToHostFloat(input, input->GetDevice().type()); + auto b_host = ToHostFloat(other, other->GetDevice().type()); + std::vector out_host(static_cast(bs * m * n)); + + for (int64_t b = 0; b < bs; ++b) { + HostGemm(false, false, m, n, k, a_host.data() + b * m * k, k, b_host.data() + b * k * n, n, 0.0f, + out_host.data() + b * m * n, n); + } + + FromHostFloat(out_host, output, output->GetDevice().type()); + return output; +} + +std::shared_ptr MatmulBackwardInput(const std::shared_ptr &other, + const std::shared_ptr &grad_output, + const std::vector &input_dims) { + const auto &other_dims = other->Dims(); + const auto &grad_dims = grad_output->Dims(); + const int64_t m = grad_dims[grad_dims.size() - 2]; + const int64_t n = grad_dims[grad_dims.size() - 1]; + const int64_t k = other_dims[other_dims.size() - 2]; + const int64_t bs = NumElementsOf(std::vector(grad_dims.begin(), grad_dims.end() - 2)); + + auto grad_input = std::make_shared(input_dims, grad_output->Dtype(), grad_output->GetDevice()); + auto g_host = ToHostFloat(grad_output, grad_output->GetDevice().type()); + auto b_host = ToHostFloat(other, other->GetDevice().type()); + std::vector gi_host(static_cast(bs * m * k)); + + // grad_input(bs,m,k) = grad(bs,m,n) * other(bs,k,n)^T + for (int64_t b = 0; b < bs; ++b) { + HostGemm(false, true, m, k, n, g_host.data() + b * m * n, n, b_host.data() + b * k * n, n, 0.0f, + gi_host.data() + b * m * k, k); + } + + FromHostFloat(gi_host, grad_input, grad_input->GetDevice().type()); + return grad_input; +} + +std::shared_ptr MatmulBackwardOther(const std::shared_ptr &input1, + const std::shared_ptr &grad_output, + const std::vector &other_dims) { + const auto &input_dims = input1->Dims(); + const auto &grad_dims = grad_output->Dims(); + 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_dims[grad_dims.size() - 1]; + const int64_t bs = NumElementsOf(std::vector(input_dims.begin(), input_dims.end() - 2)); + + auto grad_other = std::make_shared(other_dims, grad_output->Dtype(), grad_output->GetDevice()); + auto a_host = ToHostFloat(input1, input1->GetDevice().type()); + auto g_host = ToHostFloat(grad_output, grad_output->GetDevice().type()); + std::vector go_host(static_cast(bs * k * n)); + + // grad_other(bs,k,n) = input(bs,m,k)^T * grad(bs,m,n) + for (int64_t b = 0; b < bs; ++b) { + HostGemm(true, false, k, n, m, a_host.data() + b * m * k, k, g_host.data() + b * m * n, n, 0.0f, + go_host.data() + b * k * n, n); + } + + FromHostFloat(go_host, grad_other, grad_other->GetDevice().type()); + return grad_other; +} + +} // namespace infini_train::kernels::mlu + +#define REGISTER_MLU_LINEAR_KERNEL(kernel_name) \ + REGISTER_KERNEL(infini_train::Device::DeviceType::kPrivateUse1, kernel_name, \ + infini_train::kernels::mlu::kernel_name) + +REGISTER_MLU_LINEAR_KERNEL(LinearForward) +REGISTER_MLU_LINEAR_KERNEL(LinearBackwardInput) +REGISTER_MLU_LINEAR_KERNEL(LinearBackwardWeight) +REGISTER_MLU_LINEAR_KERNEL(LinearBackwardBias) +REGISTER_MLU_LINEAR_KERNEL(MatmulForward) +REGISTER_MLU_LINEAR_KERNEL(MatmulBackwardInput) +REGISTER_MLU_LINEAR_KERNEL(MatmulBackwardOther) + +#undef REGISTER_MLU_LINEAR_KERNEL diff --git a/backends/mlu/src/kernels/mlu_kernel_api.h b/backends/mlu/src/kernels/mlu_kernel_api.h new file mode 100644 index 0000000..2d10809 --- /dev/null +++ b/backends/mlu/src/kernels/mlu_kernel_api.h @@ -0,0 +1,193 @@ +#pragma once + +// MLU kernel 启动器接口头。 +// +// 关键约束：cncc 基于 clang 11（C++17 级别），编译不了框架的 C++20 头文件 +// （std::format、模板 lambda 等），所以 .mlu 文件只允许包含本头 + cnrt/bang +// 头，保持 C++11 干净；框架集成（dispatch、REGISTER_KERNEL）全部放在由 g++ +// 编译的同名 .cc 里，通过这里声明的 *Launch 函数调起设备 kernel。 + +#include + +#include + +namespace infini_train::kernels::mlu { + +// 与 mlu/runtime/mlu_dispatch.h 的 MluTypeTag() 一一对应 +enum class MluScalarType : int { + kBool = 0, + kU8, + kI8, + kU16, + kI16, + kU32, + kI32, + kU64, + kI64, + kF16, + kBF16, + kF32, + kF64, +}; + +// 填充：value_bits 按 type 位宽取低字节解释 +void FillLaunch(void *data, uint64_t value_bits, MluScalarType type, int64_t num_elements, cnrtQueue_t queue); + +void CastLaunch(void *dst, const void *src, MluScalarType dst_type, MluScalarType src_type, int64_t num_elements, + cnrtQueue_t queue); + +// 梯度累加：tensor += grad * rate +void AccumulateGradLaunch(const void *grad, float rate, void *tensor, MluScalarType type, int64_t num_elements, + cnrtQueue_t queue); + +// Adam 更新（bias correction 在 host 侧算好传入） +void AdamAccumulateGradLaunch(const void *grad, void *param, void *m, void *v, float learning_rate, float beta1, + float beta2, float eps, float bias_correction_m, float bias_correction_v, + MluScalarType type, int64_t num_elements, cnrtQueue_t queue); + +// ----------------------------------------------------------------------------- +// elementwise +// ----------------------------------------------------------------------------- + +// Broadcast 元数据（POD，按值传入 kernel）。对应 CUDA 版 elementwise.cu 的 +// BroadcastMeta：shape 右对齐到 out 的 ndim，strides 为连续布局。 +#define MLU_MAX_BROADCAST_DIMS 8 +struct MluBroadcastMeta { + int32_t ndim; + int64_t a_strides[MLU_MAX_BROADCAST_DIMS]; + int64_t b_strides[MLU_MAX_BROADCAST_DIMS]; + int64_t out_strides[MLU_MAX_BROADCAST_DIMS]; + int64_t a_shape[MLU_MAX_BROADCAST_DIMS]; + int64_t b_shape[MLU_MAX_BROADCAST_DIMS]; +}; + +enum class MluUnaryOp : int { + kNeg = 0, + kReciprocal, + kTanh, + kRsqrt, + kExp, + kLog, + kSin, + kCos, + kSigmoid, +}; + +enum class MluBinaryOp : int { + kAdd = 0, + kSub, + kMul, + kDiv, +}; + +// 一元前向：out[i] = op(in[i])；16 位浮点在 kernel 内经 float 中转 +void UnaryForwardLaunch(void *out, const void *in, MluUnaryOp op, MluScalarType type, int64_t num_elements, + cnrtQueue_t queue); + +// 一元反向：grad_in[i] = g(grad_out[i], aux[i])；aux 按算子约定为 forward 的 +// output（tanh/exp/sigmoid）或 input（reciprocal/rsqrt/log/sin/cos） +void UnaryBackwardLaunch(void *grad_in, const void *grad_out, const void *aux, MluUnaryOp op, MluScalarType type, + int64_t num_elements, cnrtQueue_t queue); + +// Pow：scalar_is_base=false 时 out=in^exponent；否则 out=exponent^in +void PowForwardLaunch(void *out, const void *in, float exponent, bool scalar_is_base, MluScalarType type, + int64_t num_elements, cnrtQueue_t queue); +void PowBackwardLaunch(void *grad_in, const void *grad_out, const void *in, float exponent, bool scalar_is_base, + MluScalarType type, int64_t num_elements, cnrtQueue_t queue); + +// 二元前向（广播）：out[i] = op(a[bcast_a(i)], b[bcast_b(i)]) +void BinaryForwardLaunch(void *out, const void *a, const void *b, MluBinaryOp op, MluScalarType type, + MluBroadcastMeta meta, int64_t num_elements, cnrtQueue_t queue); + +// 二元反向：同形状时一次扫描直接写 grad_a/grad_b；有广播时对每个 grad 元素 +// 串行 gather 累加 collapsed 维（FP32 累加，确定性顺序，不用 atomic） +void BinaryBackwardLaunch(void *grad_a, void *grad_b, const void *grad_out, const void *a, const void *b, + MluBinaryOp op, MluScalarType type, MluBroadcastMeta meta, int64_t numel_a, int64_t numel_b, + int64_t numel_out, cnrtQueue_t queue); + +// 标量二元（仅 Add/Mul；Sub/Div 标量在框架层已折合为 Add/Mul） +void ScalarForwardLaunch(void *out, const void *in, float scalar, MluBinaryOp op, MluScalarType type, + int64_t num_elements, cnrtQueue_t queue); +void ScalarBackwardLaunch(void *grad_in, const void *grad_out, float scalar, MluBinaryOp op, MluScalarType type, + int64_t num_elements, cnrtQueue_t queue); + +// 相等比较（输出 bool）；用于 causal mask 生成等。反向无意义（框架层 FATAL） +void EqualsForwardLaunch(void *out, const void *a, const void *b, MluScalarType type, MluBroadcastMeta meta, + int64_t num_elements, cnrtQueue_t queue); +void EqualsScalarForwardLaunch(void *out, const void *in, double scalar, MluScalarType type, int64_t num_elements, + cnrtQueue_t queue); + +// ----------------------------------------------------------------------------- +// embedding +// ----------------------------------------------------------------------------- + +// input 为 int64 token ids；越界 id 直接跳过（与 CUDA 版一致） +void EmbeddingForwardLaunch(const void *input, void *output, const void *weight, int64_t num_tokens, int64_t embed_dim, + int64_t vocab_size, MluScalarType type, cnrtQueue_t queue); +// grad_weight 调用前必须已清零（kernel 内为 atomic 累加） +void EmbeddingBackwardLaunch(const void *input, const void *grad_output, void *grad_weight, int64_t num_tokens, + int64_t embed_dim, int64_t vocab_size, MluScalarType type, cnrtQueue_t queue); + +// ----------------------------------------------------------------------------- +// layernorm / softmax / cross_entropy（行/样本级归约，全程 FP32 中间精度） +// ----------------------------------------------------------------------------- + +// input: (num_tokens, embed_dim)；mean_out/rstd_out 为 FP32，必须非空 +void LayerNormForwardLaunch(const void *input, const void *weight, const void *bias, float *mean_out, float *rstd_out, + void *output, float eps, int64_t num_tokens, int64_t embed_dim, MluScalarType type, + cnrtQueue_t queue); +void LayerNormBackwardLaunch(const void *input, const void *grad_output, const float *mean, const float *rstd, + const void *weight, void *grad_input, void *grad_weight, void *grad_bias, + int64_t num_tokens, int64_t embed_dim, MluScalarType type, cnrtQueue_t queue); + +// 沿 dim 的 softmax：outer/axis/inner 由 host 按 dims 拆解 +void SoftmaxForwardLaunch(void *output, const void *input, int64_t outer_size, int64_t axis_size, int64_t inner_size, + MluScalarType type, cnrtQueue_t queue); +void SoftmaxBackwardLaunch(void *grad_input, const void *grad_output, const void *output, int64_t outer_size, + int64_t axis_size, int64_t inner_size, MluScalarType type, cnrtQueue_t queue); + +// target_type ∈ {kU8, kI64}；input_type 为浮点。loss 为每样本值（bs,）， +// 批均值由 host 侧 FP64 累加完成（与 CUDA 版一致） +void CrossEntropyForwardLaunch(const void *input, const void *target, void *loss, int64_t bs, int64_t num_classes, + MluScalarType target_type, MluScalarType input_type, cnrtQueue_t queue); +void CrossEntropyBackwardLaunch(const void *input, void *grad_input, const void *target, const void *grad_output, + int64_t bs, int64_t num_classes, MluScalarType target_type, MluScalarType input_type, + cnrtQueue_t queue); + +// ----------------------------------------------------------------------------- +// transform（transpose / slice / tril / triu / mask） +// ----------------------------------------------------------------------------- + +// 通用索引元数据（POD，按值传入 kernel；对应 CUDA 版的 device 端 dims/strides 缓冲） +#define MLU_MAX_META_DIMS 8 +struct MluIndexMeta { + int32_t ndim; + int64_t dims[MLU_MAX_META_DIMS]; // slice: new_dims（输出形状） + int64_t in_strides[MLU_MAX_META_DIMS]; + int64_t out_strides[MLU_MAX_META_DIMS]; + int64_t starts[MLU_MAX_META_DIMS]; // 仅 slice 使用 + int64_t steps[MLU_MAX_META_DIMS]; // 仅 slice 使用 +}; + +// Tril/Triu 前向反向同形（保留或置 0），upper=false 为 tril +void TrilTriuLaunch(const void *input, void *output, int64_t rows, int64_t cols, int64_t diagonal, bool upper, + MluScalarType type, cnrtQueue_t queue); + +// meta.in_strides 为输入步长，meta.out_strides 为输出步长；反向即 dim0/dim1 互换再调一次 +void TransposeLaunch(const void *input, void *output, MluIndexMeta meta, int64_t dim0, int64_t dim1, + int64_t num_elements, MluScalarType type, cnrtQueue_t queue); + +// meta 需填 dims(=new_dims)/in_strides/out_strides/starts/steps；反向的 grad_input 调用前必须已清零 +void SliceForwardLaunch(const void *input, void *output, MluIndexMeta meta, int64_t num_elements, MluScalarType type, + cnrtQueue_t queue); +void SliceBackwardLaunch(const void *grad_output, void *grad_input, MluIndexMeta meta, int64_t num_elements, + MluScalarType type, cnrtQueue_t queue); + +// MaskedFill：lead=true 时 mask 覆盖前导维（mask[i/inner]），否则尾部（mask[i%mask_extent]）； +// mask 在 host 侧已 Cast 成与 input 相同 dtype，按 == 1 判定 +void MaskForwardLaunch(const void *input, const void *mask, void *output, float value, bool lead, int64_t mask_extent, + int64_t inner, int64_t num_elements, MluScalarType type, cnrtQueue_t queue); +void MaskBackwardLaunch(const void *grad_output, const void *mask, void *grad_input, bool lead, int64_t mask_extent, + int64_t inner, int64_t num_elements, MluScalarType type, cnrtQueue_t queue); + +} // namespace infini_train::kernels::mlu diff --git a/backends/mlu/src/kernels/no_op.cc b/backends/mlu/src/kernels/no_op.cc new file mode 100644 index 0000000..ccf521a --- /dev/null +++ b/backends/mlu/src/kernels/no_op.cc @@ -0,0 +1,42 @@ +// 对应上游 src/kernels/cpu/no_op.cc（NoOp 不涉及设备计算，纯 view 操作， +// 各后端实现相同）。 + +#include + +#include +#include +#include +#include + +#include "glog/logging.h" + +#include "infini_train/include/device.h" +#include "infini_train/include/dispatcher.h" +#include "infini_train/include/tensor.h" + +namespace infini_train::kernels::mlu { +std::shared_ptr NoOpForward(const std::shared_ptr &input, const std::vector &dims) { + const int64_t num_elements = std::accumulate(dims.begin(), dims.end(), 1, std::multiplies()); + CHECK_EQ(input->NumElements(), num_elements); + + auto output = std::make_shared(*input, 0, dims); + return output; +} + +std::shared_ptr NoOpBackward(const std::vector &dims, const std::shared_ptr &grad_output) { + auto num_elements = std::accumulate(dims.begin(), dims.end(), 1, std::multiplies()); + CHECK_EQ(num_elements, grad_output->NumElements()); + + auto grad_input = std::make_shared(*grad_output, 0, dims); + return grad_input; +} +} // namespace infini_train::kernels::mlu + +#define REGISTER_MLU_NO_OP_KERNEL(kernel_name) \ + REGISTER_KERNEL(infini_train::Device::DeviceType::kPrivateUse1, kernel_name, \ + infini_train::kernels::mlu::kernel_name) + +REGISTER_MLU_NO_OP_KERNEL(NoOpForward) +REGISTER_MLU_NO_OP_KERNEL(NoOpBackward) + +#undef REGISTER_MLU_NO_OP_KERNEL diff --git a/backends/mlu/src/kernels/softmax.cc b/backends/mlu/src/kernels/softmax.cc new file mode 100644 index 0000000..46109d1 --- /dev/null +++ b/backends/mlu/src/kernels/softmax.cc @@ -0,0 +1,96 @@ +// 对应上游 src/kernels/cuda/softmax.cu 的 host 集成部分；设备 kernel 在 softmax.mlu。 + +#include +#include + +#include "glog/logging.h" + +#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 "kernels/mlu_kernel_api.h" +#include "runtime/mlu_debug.h" +#include "runtime/mlu_dispatch.h" +#include "runtime/mlu_runtime_common.h" + +namespace infini_train::kernels::mlu { +namespace { +cnrtQueue_t CurrentQueue(const Device &device) { + return dynamic_cast( + infini_train::core::GetDeviceGuardImpl(device.type())->GetStream(device)) + ->mlu_queue(); +} + +void ComputeSoftmaxSizes(const std::vector &dims, int64_t dim, int64_t *outer_size, int64_t *axis_size, + int64_t *inner_size) { + *outer_size = 1; + *axis_size = dims[dim]; + *inner_size = 1; + for (int64_t i = 0; i < dim; ++i) { *outer_size *= dims[i]; } + for (size_t i = dim + 1; i < dims.size(); ++i) { *inner_size *= dims[i]; } +} +} // namespace + +std::shared_ptr SoftmaxForward(const std::shared_ptr &input, int64_t dim) { + const auto &input_dims = input->Dims(); + dim = dim < 0 ? dim + input_dims.size() : dim; + CHECK(dim >= 0 && dim < static_cast(input_dims.size())); + + auto output = std::make_shared(input_dims, input->Dtype(), input->GetDevice()); + + int64_t outer_size, axis_size, inner_size; + ComputeSoftmaxSizes(input_dims, dim, &outer_size, &axis_size, &inner_size); + if (axis_size == 0 || outer_size == 0) { + return output; + } + + const auto queue = CurrentQueue(input->GetDevice()); + core::mlu::DispatchMluFunc( + input->Dtype(), + [=]() { + SoftmaxForwardLaunch(output->DataPtr(), input->DataPtr(), outer_size, axis_size, inner_size, + MluTypeTag(), queue); + ::infini_train::core::mlu::DebugSyncQueue(queue, "SoftmaxForward"); + ::infini_train::core::mlu::DebugCheck(output, "SoftmaxForward"); + }, + "MLU SoftmaxForward"); + return output; +} + +std::shared_ptr SoftmaxBackward(const std::shared_ptr &grad_output, + const std::shared_ptr &output, int64_t dim) { + const auto &output_dims = output->Dims(); + dim = dim < 0 ? dim + output_dims.size() : dim; + CHECK(dim >= 0 && dim < static_cast(output_dims.size())); + + auto grad_input = std::make_shared(output_dims, output->Dtype(), output->GetDevice()); + + int64_t outer_size, axis_size, inner_size; + ComputeSoftmaxSizes(output_dims, dim, &outer_size, &axis_size, &inner_size); + if (axis_size == 0 || outer_size == 0) { + return grad_input; + } + + const auto queue = CurrentQueue(output->GetDevice()); + core::mlu::DispatchMluFunc( + output->Dtype(), + [=]() { + SoftmaxBackwardLaunch(grad_input->DataPtr(), grad_output->DataPtr(), output->DataPtr(), outer_size, + axis_size, inner_size, MluTypeTag(), queue); + ::infini_train::core::mlu::DebugSyncQueue(queue, "SoftmaxBackward"); + }, + "MLU SoftmaxBackward"); + return grad_input; +} +} // namespace infini_train::kernels::mlu + +#define REGISTER_MLU_SOFTMAX_KERNEL(kernel_name) \ + REGISTER_KERNEL(infini_train::Device::DeviceType::kPrivateUse1, kernel_name, \ + infini_train::kernels::mlu::kernel_name) + +REGISTER_MLU_SOFTMAX_KERNEL(SoftmaxForward) +REGISTER_MLU_SOFTMAX_KERNEL(SoftmaxBackward) + +#undef REGISTER_MLU_SOFTMAX_KERNEL diff --git a/backends/mlu/src/kernels/softmax.mlu b/backends/mlu/src/kernels/softmax.mlu new file mode 100644 index 0000000..24be8bc --- /dev/null +++ b/backends/mlu/src/kernels/softmax.mlu @@ -0,0 +1,118 @@ +// 纯 BANG C kernel + 启动器（cncc 编译，只允许 C++11 级别语法，不含框架头）。 +// 框架集成见同目录 softmax.cc。 +// 语义对应上游 src/kernels/cuda/softmax.cu（沿 dim 的 max/和归约 + 归一化）。 +// 简化策略（正确性优先）：一行一个核串行归约，全程 FP32 中间精度，顺序确定。 + +#include + +#include + +#include "kernels/mlu_kernel_api.h" + +namespace { +using infini_train::kernels::mlu::MluScalarType; + +constexpr uint32_t kDefaultTaskCount = 4; + +// 注意（实测）：本工具链下 taskIdX/coreId/taskDimX/coreDim 与预期语义不符，配合 +// cnrtFuncTypeUnion1 时循环写 GDRAM 只有首尾元素落盘。厂商示例（samples/BANG/.../ +// vectoradd/0_scalarAdd_single_core）用的是纯串行循环 + cnrtFuncTypeBlock。 +// 因此所有 kernel 统一改为单任务串行；性能优化属任务三，届时再按官方多核范式改写。 + +// 行 = (group, inner_idx) 组合，共 outer_size * inner_size 行 +template +__mlu_global__ void SoftmaxForwardKernel(T *output, const T *input, int64_t outer_size, int64_t axis_size, + int64_t inner_size) { + const int64_t rows = outer_size * inner_size; + for (int64_t row = 0; row < rows; ++row) { + const int64_t group = row / inner_size; + const int64_t inner_idx = row % inner_size; + const int64_t base = group * axis_size * inner_size + inner_idx; + + float row_max = -__builtin_inff(); + for (int64_t axis = 0; axis < axis_size; ++axis) { + const float v = static_cast(input[base + axis * inner_size]); + row_max = v > row_max ? v : row_max; + } + float row_sum = 0.0f; + for (int64_t axis = 0; axis < axis_size; ++axis) { + const int64_t idx = base + axis * inner_size; + const float e = expf(static_cast(input[idx]) - row_max); + output[idx] = static_cast(e); + row_sum += e; + } + for (int64_t axis = 0; axis < axis_size; ++axis) { + const int64_t idx = base + axis * inner_size; + output[idx] = static_cast(static_cast(output[idx]) / row_sum); + } + } +} + +template +__mlu_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) { + const int64_t rows = outer_size * inner_size; + for (int64_t row = 0; row < rows; ++row) { + const int64_t group = row / inner_size; + const int64_t inner_idx = row % inner_size; + const int64_t base = group * axis_size * inner_size + inner_idx; + + float row_sum = 0.0f; + for (int64_t axis = 0; axis < axis_size; ++axis) { + const int64_t idx = base + axis * inner_size; + row_sum += static_cast(grad_output[idx]) * static_cast(output[idx]); + } + for (int64_t axis = 0; axis < axis_size; ++axis) { + const int64_t idx = base + axis * inner_size; + grad_input[idx] + = static_cast(static_cast(output[idx]) * (static_cast(grad_output[idx]) - row_sum)); + } + } +} + +template +void SoftmaxForwardTyped(void *output, const void *input, int64_t outer_size, int64_t axis_size, int64_t inner_size, + cnrtQueue_t queue) { + const cnrtDim3_t dim = {1, 1, 1}; + SoftmaxForwardKernel<<>>(static_cast(output), static_cast(input), + outer_size, axis_size, inner_size); +} + +template +void SoftmaxBackwardTyped(void *grad_input, const void *grad_output, const void *output, int64_t outer_size, + int64_t axis_size, int64_t inner_size, cnrtQueue_t queue) { + const cnrtDim3_t dim = {1, 1, 1}; + SoftmaxBackwardKernel + <<>>(static_cast(grad_input), static_cast(grad_output), + static_cast(output), outer_size, axis_size, inner_size); +} + +#define MLU_SM_FLOAT_SWITCH(fn, ...) \ + switch (type) { \ + case MluScalarType::kF16: \ + fn(__VA_ARGS__); \ + return; \ + case MluScalarType::kBF16: \ + fn(__VA_ARGS__); \ + return; \ + case MluScalarType::kF32: \ + fn(__VA_ARGS__); \ + return; \ + default: \ + return; \ + } +} // namespace + +namespace infini_train::kernels::mlu { + +void SoftmaxForwardLaunch(void *output, const void *input, int64_t outer_size, int64_t axis_size, int64_t inner_size, + MluScalarType type, cnrtQueue_t queue) { + MLU_SM_FLOAT_SWITCH(SoftmaxForwardTyped, output, input, outer_size, axis_size, inner_size, queue) +} + +void SoftmaxBackwardLaunch(void *grad_input, const void *grad_output, const void *output, int64_t outer_size, + int64_t axis_size, int64_t inner_size, MluScalarType type, cnrtQueue_t queue) { + MLU_SM_FLOAT_SWITCH(SoftmaxBackwardTyped, grad_input, grad_output, output, outer_size, axis_size, inner_size, queue) +} + +} // namespace infini_train::kernels::mlu diff --git a/backends/mlu/src/kernels/transform.cc b/backends/mlu/src/kernels/transform.cc new file mode 100644 index 0000000..e998448 --- /dev/null +++ b/backends/mlu/src/kernels/transform.cc @@ -0,0 +1,294 @@ +// 对应上游 src/kernels/cuda/transform.cu 与 slice.cu 的 host 集成部分；设备 kernel 在 transform.mlu。 +// 只保留 GPT-2 训练路径需要的算子：Tril/Triu、Transpose、Mask（MaskedFill）、Slice。 + +#include +#include +#include +#include + +#include "glog/logging.h" + +#include "infini_train/include/common/common.h" +#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 "kernels/mlu_kernel_api.h" +#include "runtime/mlu_debug.h" +#include "runtime/mlu_dispatch.h" +#include "runtime/mlu_runtime_common.h" + +namespace infini_train::kernels::mlu { +namespace { + +cnrtQueue_t CurrentQueue(const Device &device) { + return dynamic_cast( + infini_train::core::GetDeviceGuardImpl(device.type())->GetStream(device)) + ->mlu_queue(); +} + +// 与 CUDA 版 transform.cu 的 IsLeadMaskShape / IsTailMaskShape 一致 +bool IsLeadMaskShape(const std::vector &in, const std::vector &mk) { + if (mk.empty() || in.empty() || 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; +} + +bool IsTailMaskShape(const std::vector &in, const std::vector &mk) { + if (mk.size() > in.size()) { + return false; + } + const size_t k = mk.size(); + for (size_t i = 0; i < k; ++i) { + if (!(mk[i] == in[in.size() - k + i] || mk[i] == 1)) { + return false; + } + } + return true; +} + +// 默认优先 tail（与 CUDA 版 DecideMaskMode 一致） +bool DecideMaskLead(const std::vector &in, const std::vector &mk) { + const bool lead = IsLeadMaskShape(in, mk); + const bool tail = IsTailMaskShape(in, mk); + CHECK(lead || tail) << "Mask must align/broadcast to either leading or trailing axes."; + return !tail; +} + +MluIndexMeta MakeSliceMeta(const std::vector &in_dims, const std::vector &new_dims, + const std::vector &starts, const std::vector &steps) { + CHECK_LE(in_dims.size(), MLU_MAX_META_DIMS); + MluIndexMeta meta{}; + meta.ndim = static_cast(in_dims.size()); + const auto in_strides = ComputeStrides(in_dims); + const auto out_strides = ComputeStrides(new_dims); + for (size_t i = 0; i < in_dims.size(); ++i) { + meta.dims[i] = new_dims[i]; + meta.in_strides[i] = in_strides[i]; + meta.out_strides[i] = out_strides[i]; + meta.starts[i] = starts[i]; + meta.steps[i] = steps[i]; + } + return meta; +} + +std::vector SliceNewDims(const std::vector &starts, const std::vector &ends, + const std::vector &steps) { + std::vector new_dims; + for (size_t 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]); + } + return new_dims; +} + +} // namespace + +// ----------------------------- Tril / Triu ----------------------------- + +std::shared_ptr TrilForward(const std::shared_ptr &input, int64_t diagonal) { + CHECK_EQ(input->Dims().size(), 2); + auto output = std::make_shared(input->Dims(), input->Dtype(), input->GetDevice()); + const auto queue = CurrentQueue(input->GetDevice()); + core::mlu::DispatchMluFunc( + input->Dtype(), + [=]() { + TrilTriuLaunch(input->DataPtr(), output->DataPtr(), input->Dims()[0], input->Dims()[1], diagonal, false, + MluTypeTag(), queue); + ::infini_train::core::mlu::DebugSyncQueue(queue, "TrilForward"); + }, + "MLU TrilForward"); + return output; +} + +std::shared_ptr TrilBackward(const std::shared_ptr &grad_output, int64_t diagonal) { + return TrilForward(grad_output, diagonal); +} + +std::shared_ptr TriuForward(const std::shared_ptr &input, int64_t diagonal) { + CHECK_EQ(input->Dims().size(), 2); + auto output = std::make_shared(input->Dims(), input->Dtype(), input->GetDevice()); + const auto queue = CurrentQueue(input->GetDevice()); + core::mlu::DispatchMluFunc( + input->Dtype(), + [=]() { + TrilTriuLaunch(input->DataPtr(), output->DataPtr(), input->Dims()[0], input->Dims()[1], diagonal, true, + MluTypeTag(), queue); + ::infini_train::core::mlu::DebugSyncQueue(queue, "TriuForward"); + }, + "MLU TriuForward"); + return output; +} + +std::shared_ptr TriuBackward(const std::shared_ptr &grad_output, int64_t diagonal) { + return TriuForward(grad_output, diagonal); +} + +// ----------------------------- Transpose ----------------------------- + +std::shared_ptr TransposeForward(const std::shared_ptr &input, int64_t dim0, int64_t dim1) { + const auto in_dims = input->Dims(); + CHECK_LE(in_dims.size(), MLU_MAX_META_DIMS); + const int64_t ndim = in_dims.size(); + dim0 = dim0 < 0 ? dim0 + ndim : dim0; + dim1 = dim1 < 0 ? dim1 + ndim : dim1; + CHECK(dim0 >= 0 && dim0 < ndim && dim1 >= 0 && dim1 < ndim); + + auto out_dims = in_dims; + std::swap(out_dims[dim0], out_dims[dim1]); + auto output = std::make_shared(out_dims, input->Dtype(), input->GetDevice()); + + MluIndexMeta meta{}; + meta.ndim = static_cast(ndim); + const auto in_strides = ComputeStrides(in_dims); + const auto out_strides = ComputeStrides(out_dims); + for (int64_t i = 0; i < ndim; ++i) { + meta.in_strides[i] = in_strides[i]; + meta.out_strides[i] = out_strides[i]; + } + + const auto queue = CurrentQueue(input->GetDevice()); + core::mlu::DispatchMluFunc( + input->Dtype(), + [=]() { + TransposeLaunch(input->DataPtr(), output->DataPtr(), meta, dim0, dim1, + static_cast(input->NumElements()), MluTypeTag(), queue); + ::infini_train::core::mlu::DebugSyncQueue(queue, "TransposeForward"); + ::infini_train::core::mlu::DebugCheck(output, "TransposeForward"); + }, + "MLU TransposeForward"); + return output; +} + +// 转置是自逆的：反向即对 grad_output 做同一对维度的转置 +std::shared_ptr TransposeBackward(const std::shared_ptr &grad_output, int64_t dim0, int64_t dim1) { + return TransposeForward(grad_output, dim0, dim1); +} + +// ----------------------------- Mask（MaskedFill） ----------------------------- + +std::shared_ptr MaskForward(const std::shared_ptr &input, const std::shared_ptr &mask, + float value) { + const auto dtype = input->Dtype(); + auto mask_casted = mask->Dtype() == dtype ? mask : std::make_shared(mask->To(dtype)); + CHECK_EQ(static_cast(dtype), static_cast(mask_casted->Dtype())) + << "For now, input/mask dtypes must match."; + + const bool lead = DecideMaskLead(input->Dims(), mask_casted->Dims()); + auto output = std::make_shared(input->Dims(), dtype, input->GetDevice()); + + const int64_t mask_extent = mask_casted->NumElements(); + const int64_t num_elements = input->NumElements(); + const int64_t inner = num_elements / mask_extent; + const auto queue = CurrentQueue(input->GetDevice()); + core::mlu::DispatchMluFunc( + dtype, + [=]() { + MaskForwardLaunch(input->DataPtr(), mask_casted->DataPtr(), output->DataPtr(), value, lead, mask_extent, + inner, num_elements, MluTypeTag(), queue); + ::infini_train::core::mlu::DebugSyncQueue(queue, "MaskForward"); + ::infini_train::core::mlu::DebugCheck(output, "MaskForward"); + }, + "MLU MaskForward"); + return output; +} + +std::shared_ptr MaskBackward(const std::shared_ptr &grad_output, const std::shared_ptr &mask) { + const auto dtype = grad_output->Dtype(); + auto mask_casted = mask->Dtype() == dtype ? mask : std::make_shared(mask->To(dtype)); + + const bool lead = DecideMaskLead(grad_output->Dims(), mask_casted->Dims()); + auto grad_input = std::make_shared(grad_output->Dims(), dtype, grad_output->GetDevice()); + + const int64_t mask_extent = mask_casted->NumElements(); + const int64_t num_elements = grad_output->NumElements(); + const int64_t inner = num_elements / mask_extent; + const auto queue = CurrentQueue(grad_output->GetDevice()); + core::mlu::DispatchMluFunc( + dtype, + [=]() { + MaskBackwardLaunch(grad_output->DataPtr(), mask_casted->DataPtr(), grad_input->DataPtr(), lead, mask_extent, + inner, num_elements, MluTypeTag(), queue); + ::infini_train::core::mlu::DebugSyncQueue(queue, "MaskBackward"); + }, + "MLU MaskBackward"); + return grad_input; +} + +// ----------------------------- Slice ----------------------------- + +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()); + const auto &dims = input->Dims(); + CHECK_EQ(starts.size(), dims.size()); + + const auto new_dims = SliceNewDims(starts, ends, steps); + auto output = std::make_shared(new_dims, input->Dtype(), input->GetDevice()); + const auto meta = MakeSliceMeta(dims, new_dims, starts, steps); + + const auto queue = CurrentQueue(input->GetDevice()); + core::mlu::DispatchMluFunc( + input->Dtype(), + [=]() { + SliceForwardLaunch(input->DataPtr(), output->DataPtr(), meta, static_cast(output->NumElements()), + MluTypeTag(), queue); + ::infini_train::core::mlu::DebugSyncQueue(queue, "SliceForward"); + }, + "MLU SliceForward"); + return output; +} + +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()); + const auto &dims = input->Dims(); + CHECK_EQ(starts.size(), dims.size()); + + const auto new_dims = SliceNewDims(starts, ends, steps); + auto grad_input = std::make_shared(dims, grad_output->Dtype(), grad_output->GetDevice()); + // 非切片区域必须保持 0 + grad_input->Fill(0.0); + const auto meta = MakeSliceMeta(dims, new_dims, starts, steps); + + const auto queue = CurrentQueue(input->GetDevice()); + core::mlu::DispatchMluFunc( + grad_output->Dtype(), + [=]() { + SliceBackwardLaunch(grad_output->DataPtr(), grad_input->DataPtr(), meta, + static_cast(grad_output->NumElements()), MluTypeTag(), queue); + ::infini_train::core::mlu::DebugSyncQueue(queue, "SliceBackward"); + }, + "MLU SliceBackward"); + return grad_input; +} + +} // namespace infini_train::kernels::mlu + +#define REGISTER_MLU_TRANSFORM_KERNEL(kernel_name) \ + REGISTER_KERNEL(infini_train::Device::DeviceType::kPrivateUse1, kernel_name, \ + infini_train::kernels::mlu::kernel_name) + +REGISTER_MLU_TRANSFORM_KERNEL(TrilForward) +REGISTER_MLU_TRANSFORM_KERNEL(TrilBackward) +REGISTER_MLU_TRANSFORM_KERNEL(TriuForward) +REGISTER_MLU_TRANSFORM_KERNEL(TriuBackward) +REGISTER_MLU_TRANSFORM_KERNEL(TransposeForward) +REGISTER_MLU_TRANSFORM_KERNEL(TransposeBackward) +REGISTER_MLU_TRANSFORM_KERNEL(MaskForward) +REGISTER_MLU_TRANSFORM_KERNEL(MaskBackward) +REGISTER_MLU_TRANSFORM_KERNEL(SliceForward) +REGISTER_MLU_TRANSFORM_KERNEL(SliceBackward) + +#undef REGISTER_MLU_TRANSFORM_KERNEL diff --git a/backends/mlu/src/kernels/transform.mlu b/backends/mlu/src/kernels/transform.mlu new file mode 100644 index 0000000..249a419 --- /dev/null +++ b/backends/mlu/src/kernels/transform.mlu @@ -0,0 +1,246 @@ +// 纯 BANG C kernel + 启动器（cncc 编译，只允许 C++11 级别语法，不含框架头）。 +// 框架集成见同目录 transform.cc。 +// 语义对应上游 src/kernels/cuda/transform.cu 与 slice.cu。 +// 与 CUDA 版的差异：维度/步长元数据以 POD（MluIndexMeta）按值传入 kernel +// 参数空间，不再需要 cudaMallocAsync + H2D 上传。 + +#include + +#include + +#include "kernels/mlu_kernel_api.h" + +namespace { +using infini_train::kernels::mlu::MluIndexMeta; +using infini_train::kernels::mlu::MluScalarType; + +constexpr uint32_t kDefaultTaskCount = 4; + +// 注意（实测）：本工具链下 taskIdX/coreId/taskDimX/coreDim 与预期语义不符，配合 +// cnrtFuncTypeUnion1 时循环写 GDRAM 只有首尾元素落盘。厂商示例（samples/BANG/.../ +// vectoradd/0_scalarAdd_single_core）用的是纯串行循环 + cnrtFuncTypeBlock。 +// 因此所有 kernel 统一改为单任务串行；性能优化属任务三，届时再按官方多核范式改写。 + +// ----------------------------- Tril / Triu ----------------------------- +// 前向反向同形：满足条件保留原值，否则置 0 + +template +__mlu_global__ void TrilTriuKernel(const T *input, T *output, int64_t rows, int64_t cols, int64_t diagonal) { + const int64_t n = rows * cols; + for (int64_t idx = 0; idx < n; ++idx) { + const int64_t row = idx / cols; + const int64_t col = idx % cols; + const bool keep = kTriu ? (row - col + diagonal <= 0) : (row - col + diagonal >= 0); + output[idx] = keep ? input[idx] : static_cast(0); + } +} + +// ----------------------------- Transpose ----------------------------- + +template +__mlu_global__ void TransposeKernel(const T *input, T *output, MluIndexMeta meta, int64_t dim0, int64_t dim1, + int64_t num_elements) { + for (int64_t idx = 0; idx < num_elements; ++idx) { + int64_t remaining = idx; + int64_t coords[MLU_MAX_META_DIMS]; + for (int i = 0; i < meta.ndim; ++i) { + coords[i] = remaining / meta.out_strides[i]; + remaining %= meta.out_strides[i]; + } + const int64_t tmp = coords[dim0]; + coords[dim0] = coords[dim1]; + coords[dim1] = tmp; + int64_t in_flat = 0; + for (int i = 0; i < meta.ndim; ++i) { in_flat += coords[i] * meta.in_strides[i]; } + output[idx] = input[in_flat]; + } +} + +// ----------------------------- Slice ----------------------------- + +template +__mlu_global__ void SliceForwardKernel(const T *input, T *output, MluIndexMeta meta, int64_t num_elements) { + for (int64_t out_idx = 0; out_idx < num_elements; ++out_idx) { + int64_t in_index = 0; + for (int i = 0; i < meta.ndim; ++i) { + const int64_t idx = (out_idx / meta.out_strides[i]) % meta.dims[i]; + in_index += (meta.starts[i] + idx * meta.steps[i]) * meta.in_strides[i]; + } + output[out_idx] = input[in_index]; + } +} + +// grad_input 调用前必须已清零 +template +__mlu_global__ void SliceBackwardKernel(const T *grad_output, T *grad_input, MluIndexMeta meta, int64_t num_elements) { + for (int64_t out_idx = 0; out_idx < num_elements; ++out_idx) { + int64_t in_index = 0; + for (int i = 0; i < meta.ndim; ++i) { + const int64_t idx = (out_idx / meta.out_strides[i]) % meta.dims[i]; + in_index += (meta.starts[i] + idx * meta.steps[i]) * meta.in_strides[i]; + } + grad_input[in_index] = grad_output[out_idx]; + } +} + +// ----------------------------- Mask（MaskedFill） ----------------------------- +// kTail: mask[i % mask_size]；kLead: mask[i / inner] +// mask 在 host 侧已 Cast 成与 input 相同 dtype，按 == 1 判定（与 CUDA 版一致） + +template +__mlu_global__ void MaskForwardKernel(const T *input, const T *mask, T *output, T value, int64_t mask_extent, + int64_t inner, int64_t num_elements) { + for (int64_t i = 0; i < num_elements; ++i) { + const int64_t mask_idx = kLead ? (i / inner) : (i % mask_extent); + output[i] = (mask[mask_idx] == static_cast(1)) ? value : input[i]; + } +} + +template +__mlu_global__ void MaskBackwardKernel(const T *grad_output, const T *mask, T *grad_input, int64_t mask_extent, + int64_t inner, int64_t num_elements) { + for (int64_t i = 0; i < num_elements; ++i) { + const int64_t mask_idx = kLead ? (i / inner) : (i % mask_extent); + grad_input[i] = (mask[mask_idx] == static_cast(1)) ? static_cast(0) : grad_output[i]; + } +} + +// ----------------------------- 启动器 ----------------------------- + +const cnrtDim3_t kDim = {1, 1, 1}; + +template +void TrilTriuTyped(const void *input, void *output, int64_t rows, int64_t cols, int64_t diagonal, bool upper, + cnrtQueue_t queue) { + if (upper) { + TrilTriuKernel<<>>(static_cast(input), + static_cast(output), rows, cols, diagonal); + } else { + TrilTriuKernel<<>>(static_cast(input), + static_cast(output), rows, cols, diagonal); + } +} + +template +void TransposeTyped(const void *input, void *output, MluIndexMeta meta, int64_t dim0, int64_t dim1, + int64_t num_elements, cnrtQueue_t queue) { + TransposeKernel<<>>(static_cast(input), static_cast(output), + meta, dim0, dim1, num_elements); +} + +template +void SliceForwardTyped(const void *input, void *output, MluIndexMeta meta, int64_t num_elements, cnrtQueue_t queue) { + SliceForwardKernel<<>>(static_cast(input), static_cast(output), + meta, num_elements); +} + +template +void SliceBackwardTyped(const void *grad_output, void *grad_input, MluIndexMeta meta, int64_t num_elements, + cnrtQueue_t queue) { + SliceBackwardKernel<<>>(static_cast(grad_output), + static_cast(grad_input), meta, num_elements); +} + +template +void MaskForwardTyped(const void *input, const void *mask, void *output, float value, bool lead, int64_t mask_extent, + int64_t inner, int64_t num_elements, cnrtQueue_t queue) { + const T v = static_cast(value); + if (lead) { + MaskForwardKernel + <<>>(static_cast(input), static_cast(mask), + static_cast(output), v, mask_extent, inner, num_elements); + } else { + MaskForwardKernel + <<>>(static_cast(input), static_cast(mask), + static_cast(output), v, mask_extent, inner, num_elements); + } +} + +template +void MaskBackwardTyped(const void *grad_output, const void *mask, void *grad_input, bool lead, int64_t mask_extent, + int64_t inner, int64_t num_elements, cnrtQueue_t queue) { + if (lead) { + MaskBackwardKernel + <<>>(static_cast(grad_output), static_cast(mask), + static_cast(grad_input), mask_extent, inner, num_elements); + } else { + MaskBackwardKernel + <<>>(static_cast(grad_output), static_cast(mask), + static_cast(grad_input), mask_extent, inner, num_elements); + } +} + +#define MLU_TF_NUMERIC_SWITCH(fn, ...) \ + switch (type) { \ + case MluScalarType::kU8: \ + fn(__VA_ARGS__); \ + return; \ + case MluScalarType::kI8: \ + fn(__VA_ARGS__); \ + return; \ + case MluScalarType::kU16: \ + fn(__VA_ARGS__); \ + return; \ + case MluScalarType::kI16: \ + fn(__VA_ARGS__); \ + return; \ + case MluScalarType::kU32: \ + fn(__VA_ARGS__); \ + return; \ + case MluScalarType::kI32: \ + fn(__VA_ARGS__); \ + return; \ + case MluScalarType::kU64: \ + fn(__VA_ARGS__); \ + return; \ + case MluScalarType::kI64: \ + fn(__VA_ARGS__); \ + return; \ + case MluScalarType::kF16: \ + fn(__VA_ARGS__); \ + return; \ + case MluScalarType::kBF16: \ + fn(__VA_ARGS__); \ + return; \ + case MluScalarType::kF32: \ + fn(__VA_ARGS__); \ + return; \ + default: \ + return; \ + } +} // namespace + +namespace infini_train::kernels::mlu { + +void TrilTriuLaunch(const void *input, void *output, int64_t rows, int64_t cols, int64_t diagonal, bool upper, + MluScalarType type, cnrtQueue_t queue) { + MLU_TF_NUMERIC_SWITCH(TrilTriuTyped, input, output, rows, cols, diagonal, upper, queue) +} + +void TransposeLaunch(const void *input, void *output, MluIndexMeta meta, int64_t dim0, int64_t dim1, + int64_t num_elements, MluScalarType type, cnrtQueue_t queue) { + MLU_TF_NUMERIC_SWITCH(TransposeTyped, input, output, meta, dim0, dim1, num_elements, queue) +} + +void SliceForwardLaunch(const void *input, void *output, MluIndexMeta meta, int64_t num_elements, MluScalarType type, + cnrtQueue_t queue) { + MLU_TF_NUMERIC_SWITCH(SliceForwardTyped, input, output, meta, num_elements, queue) +} + +void SliceBackwardLaunch(const void *grad_output, void *grad_input, MluIndexMeta meta, int64_t num_elements, + MluScalarType type, cnrtQueue_t queue) { + MLU_TF_NUMERIC_SWITCH(SliceBackwardTyped, grad_output, grad_input, meta, num_elements, queue) +} + +void MaskForwardLaunch(const void *input, const void *mask, void *output, float value, bool lead, int64_t mask_extent, + int64_t inner, int64_t num_elements, MluScalarType type, cnrtQueue_t queue) { + MLU_TF_NUMERIC_SWITCH(MaskForwardTyped, input, mask, output, value, lead, mask_extent, inner, num_elements, queue) +} + +void MaskBackwardLaunch(const void *grad_output, const void *mask, void *grad_input, bool lead, int64_t mask_extent, + int64_t inner, int64_t num_elements, MluScalarType type, cnrtQueue_t queue) { + MLU_TF_NUMERIC_SWITCH(MaskBackwardTyped, grad_output, mask, grad_input, lead, mask_extent, inner, num_elements, + queue) +} + +} // namespace infini_train::kernels::mlu diff --git a/backends/mlu/src/runtime/mlu_debug.h b/backends/mlu/src/runtime/mlu_debug.h new file mode 100644 index 0000000..2c50978 --- /dev/null +++ b/backends/mlu/src/runtime/mlu_debug.h @@ -0,0 +1,41 @@ +#pragma once + +// 调试探针：MLU_DEBUG_CHECK=1 时，在关键 kernel 启动后把输出张量拷回 host， +// 打印 NaN/Inf 计数与数值范围，用来定位"第一个产生 NaN 的 kernel"。 +// 只在排障时开启，正常训练零开销（一次 getenv）。 +// +// 约定（001-P2 v1 F04）： +// - 按张量真实 dtype 解码，不用 sizeof(T) 猜 FP16/BF16，也不把整型当浮点解码； +// - 统计值只统计有限元素，min/max 从首个有限值初始化；无有限值时明确打印； +// - 采样有上限，日志必须能区分"全量"与"前缀采样"； +// - 启用探针时，真实 SDK 同步/拷贝失败立即终止进程；数值上出现 NaN/Inf 不终止 +// （因果 mask 的 -inf 是预期语义）。 + +#include + +#include + +#include "infini_train/include/core/runtime/device_guard.h" +#include "infini_train/include/tensor.h" + +namespace infini_train::core::mlu { + +// 开关解析：unset / "0" / "false" 关闭；"1" / "true" 打开；其他值警告后关闭。 +bool DebugFlagEnabled(const char *env_name); + +// 可单测的同步结果处理：成功返回 true；失败在探针开启时 LOG(FATAL) 终止进程。 +bool CheckSyncResult(cnrtRet_t ret, const char *kernel_name); + +void DebugSyncQueue(cnrtQueue_t queue, const char *kernel_name); + +// dtype 用 infini_train::DataType 的整数值传入（避免在纯 C 侧引入框架头）。 +void DebugCheckTensor(const void *dev_ptr, int64_t num_elements, int dtype, const char *name); + +template void DebugCheck(const std::shared_ptr &t, const char *name) { + if (!t) { + return; + } + DebugCheckTensor(t->DataPtr(), static_cast(t->NumElements()), static_cast(t->Dtype()), name); +} + +} // namespace infini_train::core::mlu diff --git a/backends/mlu/src/runtime/mlu_dispatch.h b/backends/mlu/src/runtime/mlu_dispatch.h new file mode 100644 index 0000000..f18f7c3 --- /dev/null +++ b/backends/mlu/src/runtime/mlu_dispatch.h @@ -0,0 +1,92 @@ +#pragma once + +// 对应 master 上 src/core/runtime/cuda/cuda_dispatch.h。 +// +// 本头由 host 侧（g++）编译单元使用：bang_fp16.h / bang_bf16.h 对 +// __GNUC__ 提供 host 侧 half/bfloat16_t 支持。这两个头是 vendored 拷贝 +// （third_party/bang_host/，原因见该目录 README），不是 cncc 资源目录里的。 +// 注意 .mlu 文件不要 include 本头（见 kernels/mlu_kernel_api.h 的说明）。 + +#include +#include + +#include +#include + +#include "infini_train/include/core/backend_type_map.h" +#include "infini_train/include/dtype_dispatch.h" + +#include "kernels/mlu_kernel_api.h" + +// ----------------------------------------------------------------------------- +// MLU 低精度类型映射：FP16 -> __half, BF16 -> bfloat16_t（BANG C 类型） +// 注意：host 侧 bang_fp16.h 不定义 half 别名，只能用 __half（bang_bf16.h +// 有 bfloat16_t 的 host typedef）。 +// ----------------------------------------------------------------------------- +namespace infini_train::core { +template <> struct BackendTypeMap { + using type = __half; +}; + +template <> struct BackendTypeMap { + using type = bfloat16_t; +}; +} // namespace infini_train::core + +// 注册标准（非低精度）dtype；FP16/BF16 已在上面显式特化。 +INFINI_REGISTER_STANDARD_BACKEND_TYPES(infini_train::Device::DeviceType::kPrivateUse1) + +// ----------------------------------------------------------------------------- +// MLU dispatch 类型表 +// MLU590 硬件/toolchain 不支持 FP64 设备运算（double 比较需软件运行时 +// __nedf2，cncc 后端无法选择），因此 MLU kernel dispatch 一律排除 kFLOAT64； +// 真被请求时由 dispatch 在运行期报清晰错误。 +// ----------------------------------------------------------------------------- +#define INFINI_MLU_FLOATING_TYPES DataType::kFLOAT32, DataType::kFLOAT16, DataType::kBFLOAT16 +#define INFINI_MLU_ALL_NUMERIC_TYPES INFINI_MLU_FLOATING_TYPES, INFINI_ALL_INTEGRAL_TYPES +#define INFINI_MLU_ALL_NUMERIC_AND_LOGICAL_TYPES INFINI_MLU_ALL_NUMERIC_TYPES, INFINI_LOGICAL_TYPES + +namespace infini_train::core::mlu { + +template struct MluTypeMap : BackendTypeMap {}; + +// ----------------------------------------------------------------------------- +// MLU dispatch helpers +// ----------------------------------------------------------------------------- + +template +auto DispatchMluFunc(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 DispatchMluFunc(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::mlu + +// ----------------------------------------------------------------------------- +// MluTypeTag(): C++ 类型 -> kernels::mlu::MluScalarType（给启动器用） +// ----------------------------------------------------------------------------- +namespace infini_train::kernels::mlu { + +template constexpr MluScalarType MluTypeTag(); +template <> constexpr MluScalarType MluTypeTag() { return MluScalarType::kBool; } +template <> constexpr MluScalarType MluTypeTag() { return MluScalarType::kU8; } +template <> constexpr MluScalarType MluTypeTag() { return MluScalarType::kI8; } +template <> constexpr MluScalarType MluTypeTag() { return MluScalarType::kU16; } +template <> constexpr MluScalarType MluTypeTag() { return MluScalarType::kI16; } +template <> constexpr MluScalarType MluTypeTag() { return MluScalarType::kU32; } +template <> constexpr MluScalarType MluTypeTag() { return MluScalarType::kI32; } +template <> constexpr MluScalarType MluTypeTag() { return MluScalarType::kU64; } +template <> constexpr MluScalarType MluTypeTag() { return MluScalarType::kI64; } +template <> constexpr MluScalarType MluTypeTag<__half>() { return MluScalarType::kF16; } +template <> constexpr MluScalarType MluTypeTag() { return MluScalarType::kBF16; } +template <> constexpr MluScalarType MluTypeTag() { return MluScalarType::kF32; } +template <> constexpr MluScalarType MluTypeTag() { return MluScalarType::kF64; } + +} // namespace infini_train::kernels::mlu diff --git a/backends/mlu/src/runtime/mlu_guard_impl.cc b/backends/mlu/src/runtime/mlu_guard_impl.cc new file mode 100644 index 0000000..a29b008 --- /dev/null +++ b/backends/mlu/src/runtime/mlu_guard_impl.cc @@ -0,0 +1,337 @@ +// 对应 master 上 src/core/runtime/cuda/cuda_guard_impl.cc，CUDA API 全部换成 +// CNRT/CNNL 对应物（对照表见工作区根目录 DCU_MLU_API映射速查表.md）。 +// +// 已对照服务器 /usr/local/neuware/include/cnrt.h（CNRT 7.7.3）核实： +// notifier 系列为 cnrtNotifierCreate/Destroy/Place/Wait/Query/ElapsedTime； +// queue 支持 cnrtQueueCreateWithPriority；memcpy 方向枚举是 cnrtMemTransDir_t； +// cnrtMemcpyAsync(dst, src, bytes, queue, dir) 参数顺序与 HIP 相反； +// 无 cnrtInit（自动初始化）、无 cnrtMallocAsync/cnrtFreeAsync。 + +#include "runtime/mlu_guard_impl.h" + +#include +#include +#include +#include + +#include "infini_train/include/core/runtime/runtime_common.h" +#include "infini_train/include/device.h" + +#include "common/common_mlu.h" +#include "runtime/mlu_runtime_common.h" + +namespace infini_train::core::mlu { +namespace { +constexpr int kMaxGpus = 8; +constexpr size_t kBytesPerMB = 1024ULL * 1024ULL; + +static std::array, kMaxGpus> mlu_queues; +static thread_local std::array current_queues = {}; +#ifdef USE_CNNL +// 每个线程在每个设备上独占一个 handle；handle 在线程退出后留在池里复用， +// 因为 MluBlasHandle 退出期故意不销毁底层资源。 +struct BlasHandlePool { + std::mutex mutex; + std::array>, kMaxGpus> handles; +}; +static BlasHandlePool mlu_blas_handle_pool; +static thread_local std::array, kMaxGpus> mlu_blas_handles; +#endif // USE_CNNL + +static std::array device_stream_flags; + +inline void CheckMluDevice(Device device) { + CHECK(device.type() == Device::DeviceType::kPrivateUse1) + << "MluGuardImpl expects PrivateUse1 (MLU) device, but got type=" << static_cast(device.type()) + << " index=" << static_cast(device.index()); + const int idx = device.index(); + CHECK(idx >= 0 && idx < kMaxGpus) << "MLU device index " << idx << " out of cache range [0, " << kMaxGpus << ")."; +} + +inline cnrtNotifier_t GetMluNotifier(Event *event) { + auto *mlu_event = dynamic_cast(event); + CHECK_NOTNULL(mlu_event); + return mlu_event->mlu_notifier(); +} + +inline cnrtQueue_t GetMluQueue(Stream *stream) { + auto *mlu_stream = dynamic_cast(stream); + CHECK_NOTNULL(mlu_stream); + return mlu_stream->mlu_queue(); +} + +inline cnrtMemTransDir_t ToCnrtMemcpyKind(MemcpyKind kind) { + switch (kind) { + case MemcpyKind::kH2D: + return cnrtMemcpyHostToDev; + case MemcpyKind::kD2H: + return cnrtMemcpyDevToHost; + case MemcpyKind::kD2D: + return cnrtMemcpyDevToDev; + default: + LOG(FATAL) << "Invalid MemcpyKind=" << MemcpyKindToString(kind); + } + return cnrtMemcpyDevToDev; // unreachable +} +} // namespace + +void MluGuardImpl::InitSingleStream(Device device) { + CheckMluDevice(device); + + int current_device = -1; + CNRT_CHECK(cnrtGetDevice(¤t_device)); + CNRT_CHECK(cnrtSetDevice(device.index())); + + mlu_queues[device.index()] = std::make_unique(); + current_queues[device.index()] = mlu_queues[device.index()].get(); + + CNRT_CHECK(cnrtSetDevice(current_device)); +} + +void MluGuardImpl::InitSingleHandle(Device device) { +#ifdef USE_CNNL + CheckMluDevice(device); + + auto &handle = mlu_blas_handles.at(device.index()); + if (handle) { + return; + } + std::lock_guard lock(mlu_blas_handle_pool.mutex); + auto &pooled_handles = mlu_blas_handle_pool.handles.at(device.index()); + for (const auto &pooled_handle : pooled_handles) { + // use_count() == 1 表示只有池持有，即当前没有线程占用该 handle。 + if (pooled_handle.use_count() == 1) { + handle = pooled_handle; + return; + } + } + + int current_device = -1; + CNRT_CHECK(cnrtGetDevice(¤t_device)); + CNRT_CHECK(cnrtSetDevice(device.index())); + + std::call_once(device_stream_flags.at(device.index()), InitSingleStream, device); + + handle = std::make_shared(mlu_queues[device.index()].get()); + pooled_handles.push_back(handle); + + CNRT_CHECK(cnrtSetDevice(current_device)); +#endif // USE_CNNL +} + +MluGuardImpl::MluGuardImpl() {} + +// device +Device MluGuardImpl::GetDevice() const { + int current_device = -1; + CNRT_CHECK(cnrtGetDevice(¤t_device)); + return Device(Device::DeviceType::kPrivateUse1, current_device); +} + +void MluGuardImpl::SetDevice(Device device) const { + CheckMluDevice(device); + CNRT_CHECK(cnrtSetDevice(device.index())); +} + +int MluGuardImpl::DeviceCount() const { + unsigned int device_count = 0; + CNRT_CHECK(cnrtGetDeviceCount(&device_count)); + return static_cast(device_count); +} + +Device::DeviceType MluGuardImpl::Type() const { return Device::DeviceType::kPrivateUse1; } + +// stream +Stream *MluGuardImpl::GetStream(Device device) const { + CheckMluDevice(device); + std::call_once(device_stream_flags.at(device.index()), InitSingleStream, device); + if (current_queues.at(device.index()) == nullptr) { + current_queues.at(device.index()) = mlu_queues.at(device.index()).get(); + } + return current_queues.at(device.index()); +} + +Stream *MluGuardImpl::CreateStream(Device device) const { + CheckMluDevice(device); + int current_device = -1; + CNRT_CHECK(cnrtGetDevice(¤t_device)); + CNRT_CHECK(cnrtSetDevice(device.index())); + + Stream *stream = new MluStream(); + + CNRT_CHECK(cnrtSetDevice(current_device)); + return stream; +} + +Stream *MluGuardImpl::CreateStreamWithPriority(Device device, int priority) const { + CheckMluDevice(device); + int current_device = -1; + CNRT_CHECK(cnrtGetDevice(¤t_device)); + CNRT_CHECK(cnrtSetDevice(device.index())); + + Stream *stream = new MluStream(priority); + + CNRT_CHECK(cnrtSetDevice(current_device)); + return stream; +} + +void MluGuardImpl::DestroyStream(Stream *stream) const { + if (stream == nullptr) { + return; + } + auto *mlu_stream = dynamic_cast(stream); + CHECK_NOTNULL(mlu_stream); + delete mlu_stream; +} + +void MluGuardImpl::GetStreamPriorityRange(int *low, int *high) const { + // CNRT queue 无优先级概念。 + *low = 0; + *high = 0; +} + +// event +void MluGuardImpl::EventCreate(Event **event) const { *event = new MluEvent(); } + +void MluGuardImpl::EventCreateWithFlags(Event **event, EventFlag flags) const { *event = new MluEvent(flags); } + +void MluGuardImpl::EventDestroy(Event *event) const { + if (event == nullptr) { + return; + } + delete event; +} + +void MluGuardImpl::EventRecord(Event *event, Stream *stream) const { + auto notifier = GetMluNotifier(event); + auto queue = GetMluQueue(stream); + CNRT_CHECK(cnrtPlaceNotifier(notifier, queue)); +} + +void MluGuardImpl::StreamWaitEvent(Stream *stream, Event *event, uint32_t flags) const { + auto notifier = GetMluNotifier(event); + auto queue = GetMluQueue(stream); + // 注意参数顺序：CNRT 是 notifier 在前、queue 在后，与 HIP 相反。 + CNRT_CHECK(cnrtQueueWaitNotifier(notifier, queue, flags)); +} + +RuntimeStatus MluGuardImpl::EventSynchronize(Event *event) const { + auto notifier = GetMluNotifier(event); + cnrtRet_t status = cnrtWaitNotifier(notifier); + if (status == cnrtSuccess) { + return RuntimeStatus::kSuccess; + } + LOG(ERROR) << "MluGuardImpl::EventSynchronize failed: " << cnrtGetErrorStr(status); + return RuntimeStatus::kError; +} + +RuntimeStatus MluGuardImpl::EventQuery(Event *event) const { + auto notifier = GetMluNotifier(event); + cnrtRet_t status = cnrtQueryNotifier(notifier); + if (status == cnrtSuccess) { + return RuntimeStatus::kSuccess; + } + if (status == cnrtErrorNotReady) { + return RuntimeStatus::kNotReady; + } + LOG(ERROR) << "MluGuardImpl::EventQuery failed: " << cnrtGetErrorStr(status); + return RuntimeStatus::kError; +} + +float MluGuardImpl::EventElapsedTime(Event *start_event, Event *stop_event) const { + auto start_notifier = GetMluNotifier(start_event); + auto stop_notifier = GetMluNotifier(stop_event); + float elapsed_ms = 0.0f; + CNRT_CHECK(cnrtNotifierElapsedTime(start_notifier, stop_notifier, &elapsed_ms)); + return elapsed_ms; +} + +// sync +void MluGuardImpl::SynchronizeDevice(Device device) const { + auto original_device = GetDevice(); + SetDevice(device); + + CNRT_CHECK(cnrtSyncDevice()); + + SetDevice(original_device); +} + +void MluGuardImpl::SynchronizeStream(Stream *stream) const { + auto queue = GetMluQueue(stream); + CNRT_CHECK(cnrtQueueSync(queue)); +} + +// blas +BlasHandle *MluGuardImpl::GetBlasHandle(Device device) const { +#ifdef USE_CNNL + CheckMluDevice(device); + if (!mlu_blas_handles.at(device.index())) { + InitSingleHandle(device); + } + auto *handle = mlu_blas_handles.at(device.index()).get(); + // 已知坑 #1：每次取用都重新绑定当前 queue，保证与上层调度 stream 一致。 + CNNL_CHECK(cnnlSetQueue(handle->cnnl_handle(), GetMluQueue(GetStream(device)))); + return handle; +#else + LOG(FATAL) << "MluGuardImpl::GetBlasHandle requires CNNL (build with -DUSE_CNNL=ON " + "after installing the CNNL package)"; + return nullptr; // unreachable +#endif +} + +// memory +void MluGuardImpl::Malloc(void **dev_ptr, size_t size) { CNRT_CHECK(cnrtMalloc(dev_ptr, size)); } + +void MluGuardImpl::MallocAsync(void **dev_ptr, size_t size, Stream *stream) { + // 保底策略（AGENTS.md §五.6）：CNRT 7.7.3 无异步分配接口，用同步 cnrtMalloc。 + CNRT_CHECK(cnrtMalloc(dev_ptr, size)); +} + +void MluGuardImpl::Free(void *dev_ptr) { CNRT_CHECK(cnrtFree(dev_ptr)); } + +void MluGuardImpl::FreeAsync(void *dev_ptr, Stream *stream) { + // 保底策略：先等 queue 排空再同步释放，避免释放后队列中还有访问该地址的任务。 + auto queue = GetMluQueue(stream); + CNRT_CHECK(cnrtQueueSync(queue)); + CNRT_CHECK(cnrtFree(dev_ptr)); +} + +void MluGuardImpl::Memcpy(void *dst, const void *src, size_t count, MemcpyKind kind) { + // 关键：kernel 跑在设备当前 stream 上，而 cnrtMemcpy 只保证与默认流有序， + // 直接用它会在 kernel 还没跑完时就把数据拷走（实测读到上一个张量的残留）。 + // 因此走当前 stream 的异步拷贝 + 同步。 + auto queue = GetMluQueue(GetStream(GetDevice())); + if (queue != nullptr) { + CNRT_CHECK(cnrtMemcpyAsync(dst, const_cast(src), count, queue, ToCnrtMemcpyKind(kind))); + CNRT_CHECK(cnrtQueueSync(queue)); + return; + } + CNRT_CHECK(cnrtMemcpy(dst, const_cast(src), count, ToCnrtMemcpyKind(kind))); +} + +void MluGuardImpl::MemcpyAsync(void *dst, const void *src, size_t count, MemcpyKind kind, Stream *stream) { + auto queue = GetMluQueue(stream); + // 注意参数顺序：CNRT 是 queue 在前、方向在最后，与 HIP 相反。 + CNRT_CHECK(cnrtMemcpyAsync(dst, const_cast(src), count, queue, ToCnrtMemcpyKind(kind))); +} + +void MluGuardImpl::ResetMemPoolHighWatermarks(Device /*device*/) const { + // CNRT 没有 mempool 高水位概念，无可重置项。 +} + +std::pair MluGuardImpl::GetMemPoolPeakMB(Device device) const { + // CNRT 只有全局空闲/总量查询（对当前设备）：used=reserved=total-free + // （无独立显存池）。 + CheckMluDevice(device); + const auto original_device = GetDevice(); + SetDevice(device); + size_t free_bytes = 0, total_bytes = 0; + CNRT_CHECK(cnrtMemGetInfo(&free_bytes, &total_bytes)); + SetDevice(original_device); + const size_t used_mb = (total_bytes - free_bytes) / kBytesPerMB; + return {used_mb, used_mb}; +} + +INFINI_TRAIN_REGISTER_DEVICE_GUARD_IMPL(Device::DeviceType::kPrivateUse1, MluGuardImpl) + +} // namespace infini_train::core::mlu diff --git a/backends/mlu/src/runtime/mlu_guard_impl.h b/backends/mlu/src/runtime/mlu_guard_impl.h new file mode 100644 index 0000000..325d58b --- /dev/null +++ b/backends/mlu/src/runtime/mlu_guard_impl.h @@ -0,0 +1,92 @@ +#pragma once + +// 对应 master 上 src/core/runtime/cuda/cuda_guard_impl.h。 +// 注意：master 的 PrivateUse1 机制下不再新增 DeviceType::kMLU 枚举， +// MLU 后端统一占用 DeviceType::kPrivateUse1。 + +#include "infini_train/include/core/runtime/device_guard.h" +#include "infini_train/include/device.h" + +namespace infini_train::core { +class Stream; +class BlasHandle; +} // namespace infini_train::core + +namespace infini_train::core::mlu { + +class MluGuardImpl final : public DeviceGuardImpl { +public: + static void InitSingleStream(Device device); + + static void InitSingleHandle(Device device); + + MluGuardImpl(); + + // 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; + + // The pinned InfiniTrain revision (c90338c) does not declare + // DeviceGuardImpl::ExchangeStream, so this backend does not override it. The + // previous implementation was never called; if a later framework revision + // brings the interface back, implement it here and mark it override again. + + 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; +}; + +} // namespace infini_train::core::mlu diff --git a/backends/mlu/src/runtime/mlu_runtime_common.cc b/backends/mlu/src/runtime/mlu_runtime_common.cc new file mode 100644 index 0000000..03a5587 --- /dev/null +++ b/backends/mlu/src/runtime/mlu_runtime_common.cc @@ -0,0 +1,277 @@ +#include "runtime/mlu_runtime_common.h" + +#include "runtime/mlu_debug.h" + +#include +#include +#include +#include +#include + +#include "common/common_mlu.h" + +namespace infini_train::core::mlu { + +MluEvent::MluEvent(EventFlag /*flags*/) { + // CNRT notifier 没有创建 flag；EventFlag 的 timing/blocking 语义由 + // wait/query 调用点承载。 + CNRT_CHECK(cnrtNotifierCreate(¬ifier_)); +} + +MluEvent::~MluEvent() { + if (notifier_ != nullptr) { + CNRT_CHECK(cnrtNotifierDestroy(notifier_)); + } +} + +cnrtNotifier_t MluEvent::mlu_notifier() const { return notifier_; } + +MluStream::MluStream() { CNRT_CHECK(cnrtQueueCreate(&queue_)); } + +MluStream::MluStream(int priority) { CNRT_CHECK(cnrtQueueCreateWithPriority(&queue_, 0, priority)); } + +MluStream::~MluStream() { + // Do nothing. +} + +cnrtQueue_t MluStream::mlu_queue() const { return queue_; } + +#ifdef USE_CNNL +MluBlasHandle::MluBlasHandle(Stream *stream) { + CNNL_CHECK(cnnlCreate(&cnnl_handle_)); + // 已知坑 #1：handle 必须绑定当前 queue，否则异步执行顺序错乱。 + CNNL_CHECK(cnnlSetQueue(cnnl_handle_, dynamic_cast(stream)->mlu_queue())); +} + +MluBlasHandle::~MluBlasHandle() { + // Do nothing. +} + +cnnlHandle_t MluBlasHandle::cnnl_handle() const { return cnnl_handle_; } +#endif // USE_CNNL + +bool DebugFlagEnabled(const char *env_name) { + const char *raw = std::getenv(env_name); + if (raw == nullptr || raw[0] == '\0') { + return false; + } + const std::string v(raw); + if (v == "0" || v == "false" || v == "FALSE" || v == "off") { + return false; + } + if (v == "1" || v == "true" || v == "TRUE" || v == "on") { + return true; + } + LOG(WARNING) << "env " << env_name << "=" << v << " is not a recognised flag value; treated as disabled"; + return false; +} + +bool CheckSyncResult(cnrtRet_t ret, const char *kernel_name) { + if (ret == cnrtSuccess) { + return true; + } + // 探针开启时属于排障场景：SDK 失败后队列已不可信，必须停止，避免后续错误掩盖第一失败点。 + int dev_id = -1; + cnrtGetDevice(&dev_id); + LOG(FATAL) << "MLU debug: kernel '" << kernel_name << "' failed (cnrtQueueSync ret=" << static_cast(ret) + << ", device=" << dev_id << ")"; + return false; +} + +void DebugSyncQueue(cnrtQueue_t queue, const char *kernel_name) { + if (!DebugFlagEnabled("MLU_DEBUG_SYNC")) { + return; + } + CheckSyncResult(cnrtQueueSync(queue), kernel_name); + LOG(INFO) << "MLU kernel synced: " << kernel_name; +} + +void DebugCheckTensor(const void *dev_ptr, int64_t num_elements, int dtype, const char *name) { + if (!DebugFlagEnabled("MLU_DEBUG_CHECK")) { + return; + } + if (num_elements <= 0) { + return; + } + + // 探针内部也走当前流，否则读到的是流水线上的旧数据（历史踩过的坑）。 + if (cnrtSyncDevice() != cnrtSuccess) { + LOG(FATAL) << "MLU debug: cnrtSyncDevice failed before probing " << name; + } + + const int64_t kMaxSample = 262144; + const int64_t n = num_elements < kMaxSample ? num_elements : kMaxSample; + + // 按真实 dtype 解码，禁止用元素字节数猜类型 + enum : int { kBOOL = 0, kU8, kI8, kU16, kI16, kU32, kI32, kU64, kI64, kBF16, kF16, kF32, kF64 }; + int elem_size = 0; + switch (dtype) { + case kBOOL: + case kU8: + case kI8: + elem_size = 1; + break; + case kU16: + case kI16: + case kBF16: + case kF16: + elem_size = 2; + break; + case kU32: + case kI32: + case kF32: + elem_size = 4; + break; + case kU64: + case kI64: + case kF64: + elem_size = 8; + break; + default: + LOG(WARNING) << "[check] " << name << ": unsupported dtype id " << dtype << ", skipped"; + return; + } + + std::vector host(static_cast(n) * static_cast(elem_size)); + if (cnrtMemcpy(host.data(), const_cast(dev_ptr), host.size(), cnrtMemcpyDevToHost) != cnrtSuccess) { + LOG(FATAL) << "MLU debug: D2H copy failed while probing " << name; + } + + int64_t non_finite = 0; + int64_t finite = 0; + double vmin = 0.0, vmax = 0.0; + bool is_float = true; + + auto push = [&](double v, bool finite_v) { + if (!finite_v) { + ++non_finite; + return; + } + if (finite == 0) { + vmin = vmax = v; + } + vmin = v < vmin ? v : vmin; + vmax = v > vmax ? v : vmax; + ++finite; + }; + + switch (dtype) { + case kF32: { + const auto *p = reinterpret_cast(host.data()); + for (int64_t i = 0; i < n; ++i) { push(static_cast(p[i]), std::isfinite(p[i])); } + break; + } + case kBF16: { + const auto *p = reinterpret_cast(host.data()); + for (int64_t i = 0; i < n; ++i) { + const float v = infini_train::detail::Bf16BitsToFloat(p[i]); + // bf16 的 inf/nan 位型：指数位（bit7..14）全 1 + const bool finite_v = ((p[i] >> 7) & 0xFF) != 0xFF; + push(static_cast(v), finite_v); + } + break; + } + case kF16: { + const auto *p = reinterpret_cast(host.data()); + for (int64_t i = 0; i < n; ++i) { + const float v = infini_train::detail::Fp16BitsToFloat(p[i]); + // fp16 的 inf/nan 位型：指数位（bit10..14）全 1 + const bool finite_v = ((p[i] >> 10) & 0x1F) != 0x1F; + push(static_cast(v), finite_v); + } + break; + } + case kF64: { + const auto *p = reinterpret_cast(host.data()); + for (int64_t i = 0; i < n; ++i) { push(p[i], std::isfinite(p[i])); } + break; + } + case kBOOL: { + const auto *p = reinterpret_cast(host.data()); + for (int64_t i = 0; i < n; ++i) { push(static_cast(p[i]), true); } + break; + } + default: { // 整型：无 NaN 概念；按原生位宽统计，不经过 double，避免 UINT64_MAX/INT64_MAX 失真 + is_float = false; + const char *p = host.data(); + bool is_unsigned = (dtype == kU8 || dtype == kU16 || dtype == kU32 || dtype == kU64 || dtype == kBOOL); + if (is_unsigned && (dtype == kU64 || dtype == kU32 || dtype == kU16 || dtype == kU8 || dtype == kBOOL)) { + uint64_t vmin_u = 0, vmax_u = 0; + for (int64_t i = 0; i < n; ++i) { + uint64_t v = 0; + switch (dtype) { + case kBOOL: + case kU8: + v = reinterpret_cast(p)[i]; + break; + case kU16: + v = reinterpret_cast(p)[i]; + break; + case kU32: + v = reinterpret_cast(p)[i]; + break; + default: + v = reinterpret_cast(p)[i]; + break; + } + if (i == 0 || v < vmin_u) { + vmin_u = v; + } + if (i == 0 || v > vmax_u) { + vmax_u = v; + } + } + std::ostringstream os; + os << "[check] " << name << " dtype=" << dtype << " n=" << num_elements << " sampled=" << n + << (n < num_elements ? "(prefix)" : "(full)") << " non_finite=0 uint_min=" << vmin_u + << " uint_max=" << vmax_u; + LOG(INFO) << os.str(); + return; + } + int64_t vmin_s = 0, vmax_s = 0; + for (int64_t i = 0; i < n; ++i) { + int64_t v = 0; + switch (dtype) { + case kI8: + v = reinterpret_cast(p)[i]; + break; + case kI16: + v = reinterpret_cast(p)[i]; + break; + case kI32: + v = reinterpret_cast(p)[i]; + break; + default: + v = reinterpret_cast(p)[i]; + break; + } + if (i == 0 || v < vmin_s) { + vmin_s = v; + } + if (i == 0 || v > vmax_s) { + vmax_s = v; + } + } + std::ostringstream os2; + os2 << "[check] " << name << " dtype=" << dtype << " n=" << num_elements << " sampled=" << n + << (n < num_elements ? "(prefix)" : "(full)") << " non_finite=0 int_min=" << vmin_s + << " int_max=" << vmax_s; + LOG(INFO) << os2.str(); + return; + } + } + + std::ostringstream os; + os << "[check] " << name << " dtype=" << dtype << " n=" << num_elements << " sampled=" << n + << (n < num_elements ? "(prefix)" : "(full)") << " non_finite=" << non_finite; + if (finite == 0) { + os << " min=NA max=NA (no finite element)"; + } else if (is_float) { + os << " min=" << vmin << " max=" << vmax; + } else { + os << " min=" << static_cast(vmin) << " max=" << static_cast(vmax); + } + LOG(INFO) << os.str(); +} + +} // namespace infini_train::core::mlu diff --git a/backends/mlu/src/runtime/mlu_runtime_common.h b/backends/mlu/src/runtime/mlu_runtime_common.h new file mode 100644 index 0000000..ff22402 --- /dev/null +++ b/backends/mlu/src/runtime/mlu_runtime_common.h @@ -0,0 +1,68 @@ +#pragma once + +// 对应 master 上 src/core/runtime/cuda/cuda_runtime_common.h。 +// 概念映射：HIP/CUDA stream -> CNRT queue，event -> CNRT notifier， +// cublas handle -> CNNL handle。 + +#include + +#ifdef USE_CNNL +#include +#endif +#include + +#include "infini_train/include/core/runtime/runtime_common.h" + +namespace infini_train::core { +class Stream; +} + +namespace infini_train::core::mlu { + +class MluEvent final : public Event { +public: + explicit MluEvent(EventFlag flags = EventFlag::kDefault); + ~MluEvent() override; + + cnrtNotifier_t mlu_notifier() const; + +private: + cnrtNotifier_t notifier_ = nullptr; +}; + +class MluStream : public Stream { +public: + MluStream(); + explicit MluStream(int priority); + + // NOTE: 与 CUDA 实现同理——全局 stream 故意不销毁。进程退出时静态对象的 + // 析构可能发生在 CNRT runtime 销毁之后，此时调 cnrtQueueDestroy 会崩溃， + // 因此“忘掉”销毁。 + ~MluStream() override; + + cnrtQueue_t mlu_queue() const; + +private: + cnrtQueue_t queue_ = nullptr; +}; + +// 调试用：环境变量 MLU_DEBUG_SYNC=1 时在每个 kernel 启动后同步队列并打日志， +// 用来定位设备侧崩溃（CNRT 的 CN_INVOKE_ERROR 不报是哪个 kernel）。 +void DebugSyncQueue(cnrtQueue_t queue, const char *kernel_name); + +#ifdef USE_CNNL +class MluBlasHandle : public BlasHandle { +public: + explicit MluBlasHandle(Stream *stream); + + // NOTE: 同上，故意不销毁（cnnlDestroy 在退出期不安全）。 + ~MluBlasHandle() override; + + cnnlHandle_t cnnl_handle() const; + +private: + cnnlHandle_t cnnl_handle_; +}; +#endif // USE_CNNL + +} // namespace infini_train::core::mlu diff --git a/backends/mlu/tests/CMakeLists.txt b/backends/mlu/tests/CMakeLists.txt new file mode 100644 index 0000000..cfc79c8 --- /dev/null +++ b/backends/mlu/tests/CMakeLists.txt @@ -0,0 +1,19 @@ +# MLU regression tools. They only need the provider sources and the framework, +# so they are available even when the model examples are switched off. + +# Operator comparison against the CPU backend, plus the comparator self-checks +# and the controlled probe/fatal cases. +add_executable(mlu_op_test op_test.cc) +target_include_directories(mlu_op_test PRIVATE + "${CMAKE_CURRENT_SOURCE_DIR}/../src" + "${NEUWARE_HOME}/include" +) +link_infini_train_exe(mlu_op_test + EXTRA_ARCHIVES infini_train_mlu infini_train_mlu_kernels) +set_target_properties(mlu_op_test PROPERTIES + OUTPUT_NAME op_test + RUNTIME_OUTPUT_DIRECTORY "${CMAKE_BINARY_DIR}" +) + +# tests/math_selftest has no framework dependency and keeps its own +# build_and_run.sh (cncc + g++, must be run from that directory). diff --git a/backends/mlu/tests/math_selftest/build_and_run.sh b/backends/mlu/tests/math_selftest/build_and_run.sh new file mode 100644 index 0000000..8e0887e --- /dev/null +++ b/backends/mlu/tests/math_selftest/build_and_run.sh @@ -0,0 +1,24 @@ +#!/usr/bin/env bash +# 设备数学函数单测的编译 + 运行脚本（在 MLU 服务器/容器的本目录下执行） +# +# 用法：bash build_and_run.sh +# 依赖：cncc、cnrt（NEUWARE_HOME，默认 /usr/local/neuware） +# 说明：默认不带 --bang-mlu-arch（与框架 CMake 的做法一致）；需要指定时 +# MLU_ARCH=mtp_613 bash build_and_run.sh +set -euo pipefail + +NEUWARE_HOME="${NEUWARE_HOME:-/usr/local/neuware}" +ARCH="${MLU_ARCH:-}" + +echo "[1/3] cncc 编译 .mlu" +NEUWARE_HOME="$NEUWARE_HOME" "$NEUWARE_HOME/bin/cncc" -std=c++17 -fPIC -c math_selftest.mlu \ + -o math_selftest.o -I. -I"$NEUWARE_HOME/include" ${ARCH:+--bang-mlu-arch=$ARCH} + +echo "[2/3] g++ 编译 host" +g++ -std=c++17 -O2 main.cc math_selftest.o \ + -I"$NEUWARE_HOME/include" -L"$NEUWARE_HOME/lib64" -lcnrt -o math_selftest + +if [[ "${1:-}" == "--build-only" ]]; then echo "[3/3] 跳过运行（--build-only）"; exit 0; fi + +echo "[3/3] 运行" +LD_LIBRARY_PATH="$NEUWARE_HOME/lib64:${LD_LIBRARY_PATH:-}" ./math_selftest diff --git a/backends/mlu/tests/math_selftest/main.cc b/backends/mlu/tests/math_selftest/main.cc new file mode 100644 index 0000000..23c61f8 --- /dev/null +++ b/backends/mlu/tests/math_selftest/main.cc @@ -0,0 +1,320 @@ +// 设备数学函数 / 启动方式自测。 +// +// 默认套件（正确性）：单任务串行 + cnrtFuncTypeBlock 下的 exp/log/tanh/sqrt/rsqrt/pow +// 以及一次整块拷贝（identity），与 host 参考对拍；任何数值不匹配、非预期 NaN、同步失败 +// 都使进程退出码非 0。 +// +// 旧诊断模式（--legacy-union-diag）：保留历史 Union1 + 多任务索引的对照实验。 +// 该组合已知会导致循环写 GDRAM 只有首尾落盘，因此**预期失败**；退出码照实反映失败， +// 只用于回归对照，不计入默认套件的通过数。 +// +// 用法：./math_selftest [--device N] [--legacy-union-diag] + +#include +#include +#include +#include +#include +#include +#include +#include + +#include + +#define CHECK_CNRT(call) \ + do { \ + const cnrtRet_t _r = (call); \ + if (_r != cnrtSuccess) { \ + std::printf("[FATAL] CNRT failed: %s -> %d\n", #call, static_cast(_r)); \ + std::exit(3); \ + } \ + } while (0) + +namespace { +enum MathOp : int { kExp = 0, kLog = 1, kTanh = 2, kSqrt = 3, kRsqrt = 4, kPow3 = 5, kIdentity = 6, kOpCount = 7 }; + +const char *OpName(int op) { + switch (op) { + case kExp: + return "expf"; + case kLog: + return "logf"; + case kTanh: + return "tanhf"; + case kSqrt: + return "sqrtf"; + case kRsqrt: + return "rsqrt(1/sqrtf)"; + case kPow3: + return "powf(x,3)"; + case kIdentity: + return "identity(x)"; + default: + return "?"; + } +} + +float HostRef(int op, float x) { + switch (op) { + case kExp: + return std::exp(x); + case kLog: + return std::log(x); + case kTanh: + return std::tanh(x); + case kSqrt: + return std::sqrt(x); + case kRsqrt: + return 1.0f / std::sqrt(x); + case kPow3: + return x * x * x; + case kIdentity: + return x; + default: + return 0.0f; + } +} + +// 统一非有限判定（S04）：ref 非有限时要求同名同号（域内 NaN/Inf 可预期）； +// ref 有限时 got 必须有限，再比较 atol+rtol*|ref|。 +bool Matches(double a, double b, double atol, double rtol) { + const bool a_nan = std::isnan(a), b_nan = std::isnan(b); + const bool a_inf = std::isinf(a), b_inf = std::isinf(b); + if (a_nan || a_inf) { + return (a_nan && b_nan) || (a_inf && b_inf && ((a > 0) == (b > 0))); + } + if (b_nan || b_inf) { + return false; + } + return std::fabs(a - b) <= atol + rtol * std::fabs(a); +} + +// 纯 host 反例：不启动设备，验证有限/同号 Inf/异号 Inf/预期 NaN 规则 +int RunCompareSelftest() { + const double inf = std::numeric_limits::infinity(); + const double nan_v = std::numeric_limits::quiet_NaN(); + struct C { + const char *name; + double ref; + double got; + bool expect; + }; + const C cases[] = { + {"finite equal", 1.0, 1.0, true}, {"finite beyond tol", 1.0, 1.01, false}, + {"ref +Inf got finite", inf, 0.0, false}, {"ref -Inf got finite", -inf, 0.0, false}, + {"ref +Inf got -Inf", inf, -inf, false}, {"ref +Inf got +Inf", inf, inf, true}, + {"ref NaN got NaN", nan_v, nan_v, true}, {"ref NaN got finite", nan_v, 1.0, false}, + {"ref finite got NaN", 1.0, nan_v, false}, {"ref finite got Inf", 1.0, inf, false}, + }; + int bad = 0; + for (const auto &c : cases) { + const bool got = Matches(c.ref, c.got, 1e-6, 1e-5); + const bool ok = (got == c.expect); + if (!ok) { + ++bad; + } + std::cout << c.name << " expect=" << (c.expect ? "PASS" : "FAIL") << " actual=" << (got ? "PASS" : "FAIL") + << " " << (ok ? "ok" : "MISMATCH") << std::endl; + } + std::cout << "compare-selftest: " << (bad == 0 ? "PASS" : "FAIL") << " (" << bad << " mismatch)" << std::endl; + return bad == 0 ? 0 : 1; +} + +} // namespace + +void MathSelftestLaunch(const void *in, void *out, int64_t n, int op, cnrtQueue_t queue); +void Probe3Launch(const void *in, void *out, int64_t n, cnrtQueue_t queue); +void ProbeLaunch(void *out, int64_t n, cnrtQueue_t queue); +void Probe2Launch(const void *in, void *out, int64_t n, int mode, cnrtQueue_t queue); +void Probe4Launch(void *out, cnrtQueue_t queue); +void Probe5Launch(const void *in, void *out, int64_t n, cnrtQueue_t queue); + +int main(int argc, char **argv) { + int device = 0; + bool legacy = false; + for (int i = 1; i < argc; ++i) { + const std::string a(argv[i]); + if (a == "--legacy-union-diag") { + legacy = true; + } else if (a == "--compare-selftest") { + continue; // 在参数解析后、设备初始化前处理 + } else if (a == "--device" && i + 1 < argc) { + device = std::atoi(argv[++i]); + } else { + std::printf("unknown arg: %s\n", argv[i]); + return 2; + } + } + + constexpr int64_t kN = 1024; + std::vector x(kN); + for (int64_t i = 0; i < kN; ++i) { + const float t = static_cast(i) / static_cast(kN); + x[i] = -3.0f + 6.0f * t; + if (i % 64 == 0) { + x[i] = 0.5f; + } + if (i % 64 == 1) { + x[i] = 1.0f; + } + if (i % 64 == 2) { + x[i] = 2.0f; + } + if (i % 64 == 3) { + x[i] = 8.5f; + } + if (i % 64 == 4) { + x[i] = 0.001f; + } + } + + for (int i = 1; i < argc; ++i) { + if (std::string(argv[i]) == "--compare-selftest") { + return RunCompareSelftest(); + } + } + + CHECK_CNRT(cnrtSetDevice(device)); + + void *dev_in = nullptr; + void *dev_out = nullptr; + CHECK_CNRT(cnrtMalloc(&dev_in, kN * sizeof(float))); + CHECK_CNRT(cnrtMalloc(&dev_out, kN * sizeof(float))); + CHECK_CNRT(cnrtMemcpy(dev_in, x.data(), kN * sizeof(float), cnrtMemcpyHostToDev)); + + cnrtQueue_t queue = nullptr; + CHECK_CNRT(cnrtQueueCreate(&queue)); + + std::printf("math_selftest: device=%d mode=%s\n", device, legacy ? "legacy-union-diag" : "correctness(Block)"); + + int failures = 0; + + if (legacy) { + // 历史诊断：Union1 + 多任务索引，已知会丢写；退出码照实反映失败 + CHECK_CNRT(cnrtMemset(dev_out, 0, kN * sizeof(float))); + ProbeLaunch(dev_out, kN, queue); + CHECK_CNRT(cnrtQueueSync(queue)); + std::vector p(kN); + CHECK_CNRT(cnrtMemcpy(p.data(), dev_out, kN * sizeof(float), cnrtMemcpyDevToHost)); + std::printf("legacy probe(const write): out[0]=%.1f out[1]=%.1f out[n-1]=%.1f\n", p[0], p[1], p[kN - 1]); + + CHECK_CNRT(cnrtMemset(dev_out, 0, kN * sizeof(float))); + Probe2Launch(dev_in, dev_out, kN, 2, queue); + CHECK_CNRT(cnrtQueueSync(queue)); + CHECK_CNRT(cnrtMemcpy(p.data(), dev_out, kN * sizeof(float), cnrtMemcpyDevToHost)); + int64_t mismatch = 0; + for (int64_t i = 0; i < kN; ++i) { + if (p[i] != x[i]) { + ++mismatch; + } + } + std::printf("legacy probe2(read in[i]) mismatch=%lld/%lld (expect mismatch under Union1)\n", + static_cast(mismatch), static_cast(kN)); + if (mismatch != 0) { + ++failures; + } + + CHECK_CNRT(cnrtMemset(dev_out, 0, kN * sizeof(float))); + Probe4Launch(dev_out, queue); + CHECK_CNRT(cnrtQueueSync(queue)); + CHECK_CNRT(cnrtMemcpy(p.data(), dev_out, kN * sizeof(float), cnrtMemcpyDevToHost)); + std::printf("legacy probe4(const write, 1 task) out[0]=%.1f out[5]=%.1f\n", p[0], p[5]); + if (p[0] != 7.0f || p[5] != 8.0f) { + ++failures; + } + + CHECK_CNRT(cnrtMemset(dev_out, 0, kN * sizeof(float))); + Probe5Launch(dev_in, dev_out, 8, queue); + CHECK_CNRT(cnrtQueueSync(queue)); + CHECK_CNRT(cnrtMemcpy(p.data(), dev_out, 8 * sizeof(float), cnrtMemcpyDevToHost)); + int64_t small_bad = 0; + for (int64_t i = 0; i < 8; ++i) { + if (p[i] != x[i]) { + ++small_bad; + } + } + std::printf("legacy probe5(loop n=8) mismatch=%lld/8\n", static_cast(small_bad)); + if (small_bad != 0) { + ++failures; + } + + std::printf("legacy mode result: %d failing probe(s)\n", failures); + CHECK_CNRT(cnrtQueueDestroy(queue)); + cnrtFree(dev_in); + cnrtFree(dev_out); + return failures == 0 ? 0 : 1; + } + + // 默认正确性套件：Block + 串行 + { + CHECK_CNRT(cnrtMemset(dev_out, 0, kN * sizeof(float))); + Probe3Launch(dev_in, dev_out, kN, queue); + if (cnrtQueueSync(queue) != cnrtSuccess) { + // R03：队列已不可信，立即停止，不再向同一队列提交任何 kernel + std::printf("[FAIL] identity copy: queue sync failed -> abort (no further launches)\n"); + cnrtQueueDestroy(queue); + cnrtFree(dev_in); + cnrtFree(dev_out); + std::printf("RESULT: FAIL\n"); + return 3; + } else { + std::vector got(kN); + CHECK_CNRT(cnrtMemcpy(got.data(), dev_out, kN * sizeof(float), cnrtMemcpyDevToHost)); + int64_t bad = 0; + for (int64_t i = 0; i < kN; ++i) { + if (got[i] != x[i]) { + ++bad; + } + } + std::printf("%-18s %s mismatch=%lld/%lld\n", "identity (copy)", bad == 0 ? "PASS" : "FAIL", + static_cast(bad), static_cast(kN)); + if (bad != 0) { + ++failures; + } + } + } + + std::vector got(kN); + for (int op = 0; op < kOpCount; ++op) { + CHECK_CNRT(cnrtMemset(dev_out, 0, kN * sizeof(float))); + MathSelftestLaunch(dev_in, dev_out, kN, op, queue); + if (cnrtQueueSync(queue) != cnrtSuccess) { + std::printf("%-18s FAIL queue sync failed (stop)\n", OpName(op)); + ++failures; + break; // 队列已不可信，按方案停止 + } + CHECK_CNRT(cnrtMemcpy(got.data(), dev_out, kN * sizeof(float), cnrtMemcpyDevToHost)); + + double max_abs = 0.0, max_rel = 0.0; + long long bad = 0; + for (int64_t i = 0; i < kN; ++i) { + const double a = static_cast(HostRef(op, x[i])); + const double b = static_cast(got[i]); + if (!Matches(a, b, 1e-6, 1e-5)) { + ++bad; + } + if (std::isfinite(a) && std::isfinite(b)) { + const double d = std::fabs(a - b); + const double r = d / (std::fabs(a) > 1e-12 ? std::fabs(a) : 1e-12); + if (d > max_abs) { + max_abs = d; + } + if (r > max_rel) { + max_rel = r; + } + } + } + const bool ok = bad == 0; + if (!ok) { + ++failures; + } + std::printf("%-18s %s max_abs=%.3e max_rel=%.3e bad=%lld/%lld\n", OpName(op), ok ? "PASS" : "FAIL", max_abs, + max_rel, bad, static_cast(kN)); + } + + CHECK_CNRT(cnrtQueueDestroy(queue)); + cnrtFree(dev_in); + cnrtFree(dev_out); + std::printf("RESULT: %s\n", failures == 0 ? "PASS" : "FAIL"); + return failures == 0 ? 0 : 1; +} diff --git a/backends/mlu/tests/math_selftest/math_selftest.mlu b/backends/mlu/tests/math_selftest/math_selftest.mlu new file mode 100644 index 0000000..911b084 --- /dev/null +++ b/backends/mlu/tests/math_selftest/math_selftest.mlu @@ -0,0 +1,150 @@ +// 设备侧标量数学函数单测：确认哪些函数在 MLU590 上返回错误值。 +// 独立于框架，直接用 cnrt + cncc 编译运行（见同目录 build_and_run.sh）。 +// +// 背景：PowForward 曾整片 NaN → 证实设备侧 powf 不可用；softmax 输出不在 [0,1] +// → 怀疑 expf/tanhf。本测试把候选函数逐个与 host std:: 版本对拍。 + +#include + +#include + +// 与 host 侧约定的算子编号 +enum MathOp : int { + kExp = 0, + kLog = 1, + kTanh = 2, + kSqrt = 3, + kRsqrt = 4, + kPow3 = 5, + kIdentity = 6, + kOpCount = 7, +}; + +// 最小探针：不依赖任何数学函数和循环，只写常量，用来验证 launch 机制本身 +__mlu_global__ void ProbeKernel(float *out, int64_t n) { + if (taskIdX == 0 && coreId == 0) { + out[0] = 42.0f; + } + if (taskIdX == 1 && coreId == 0) { + out[1] = -7.5f; + } + out[n - 1] = 1.0f; +} + +template __mlu_func__ float Eval(float x) { + if (Op == kExp) { + return expf(x); + } + if (Op == kLog) { + return logf(x); + } + if (Op == kTanh) { + return tanhf(x); + } + if (Op == kSqrt) { + return sqrtf(x); + } + if (Op == kRsqrt) { + return 1.0f / sqrtf(x); + } + if (Op == kIdentity) { + return x; + } + return powf(x, 3.0f); +} + +template __mlu_global__ void MathKernel(const float *in, float *out, int64_t n) { + // 纯串行：不使用 taskIdX/coreId/taskDimX/coreDim（这些内建量在本工具链下 + // 与预期不符，见 probe3 对照），由调用方按 Block/单任务启动 + for (int64_t i = 0; i < n; ++i) { out[i] = Eval(in[i]); } +} + +template void LaunchTyped(const float *in, float *out, int64_t n, cnrtQueue_t queue) { + const cnrtDim3_t dim = {1, 1, 1}; // 厂商示例用 Block；Union1 下循环写不落盘 + MathKernel<<>>(in, out, n); +} + +void ProbeLaunch(void *out, int64_t n, cnrtQueue_t queue) { + const cnrtDim3_t dim = {4, 1, 1}; + ProbeKernel<<>>(static_cast(out), n); +} + +void MathSelftestLaunch(const void *in, void *out, int64_t n, int op, cnrtQueue_t queue) { + switch (op) { + case kExp: + LaunchTyped(static_cast(in), static_cast(out), n, queue); + return; + case kLog: + LaunchTyped(static_cast(in), static_cast(out), n, queue); + return; + case kTanh: + LaunchTyped(static_cast(in), static_cast(out), n, queue); + return; + case kSqrt: + LaunchTyped(static_cast(in), static_cast(out), n, queue); + return; + case kRsqrt: + LaunchTyped(static_cast(in), static_cast(out), n, queue); + return; + case kPow3: + LaunchTyped(static_cast(in), static_cast(out), n, queue); + return; + case kIdentity: + LaunchTyped(static_cast(in), static_cast(out), n, queue); + return; + default: + return; + } +} + +// 进一步切分：mode 0 = 按下标算值（不读输入）；mode 1 = 读 in[0] 广播；mode 2 = 读 in[i] +__mlu_global__ void Probe2Kernel(const float *in, float *out, int64_t n, int mode) { + const int64_t stride = static_cast(taskDimX) * coreDim; + for (int64_t i = static_cast(taskIdX) * coreDim + coreId; i < n; i += stride) { + if (mode == 0) { + out[i] = static_cast(i) * 0.5f; + } else if (mode == 1) { + out[i] = in[0]; + } else { + out[i] = in[i]; + } + } +} + +void Probe2Launch(const void *in, void *out, int64_t n, int mode, cnrtQueue_t queue) { + const cnrtDim3_t dim = {4, 1, 1}; + Probe2Kernel<<>>(static_cast(in), static_cast(out), n, + mode); +} + +// 决定性对照：单任务、不依赖 taskIdX/coreDim，纯串行写 0..n-1 +__mlu_global__ void Probe3Kernel(const float *in, float *out, int64_t n) { + for (int64_t i = 0; i < n; ++i) { out[i] = in[i]; } +} + +void Probe3Launch(const void *in, void *out, int64_t n, cnrtQueue_t queue) { + const cnrtDim3_t dim = {1, 1, 1}; + // 厂商 vectoradd 示例用的是 cnrtFuncTypeBlock + Probe3Kernel<<>>(static_cast(in), static_cast(out), n); +} + +// 微测试：单任务、纯常量写（无循环） +__mlu_global__ void Probe4Kernel(float *out) { + out[0] = 7.0f; + out[5] = 8.0f; +} + +// 微测试：单任务、单次循环（n 很小） +__mlu_global__ void Probe5Kernel(const float *in, float *out, int64_t n) { + for (int64_t i = 0; i < n; ++i) { out[i] = in[i]; } +} + +void Probe4Launch(void *out, cnrtQueue_t queue) { + const cnrtDim3_t dim = {1, 1, 1}; + Probe4Kernel<<>>(static_cast(out)); +} + +void Probe5Launch(const void *in, void *out, int64_t n, cnrtQueue_t queue) { + const cnrtDim3_t dim = {1, 1, 1}; + Probe5Kernel<<>>(static_cast(in), static_cast(out), n); +} diff --git a/backends/mlu/tests/op_test.cc b/backends/mlu/tests/op_test.cc new file mode 100644 index 0000000..911adef --- /dev/null +++ b/backends/mlu/tests/op_test.cc @@ -0,0 +1,581 @@ +// 逐算子对拍 + 比较器自检 + 探针用例（001-P2 v2 R01/R02/R05/R06）。 +// +// 判据（R01 修正后）： +// 1) 先按显式期望分类：ref 为 NaN/±Inf 时要求 got 同类（NaN↔NaN、同号 Inf）， +// 且只有声明"允许特殊参考"的域用例才接受 ref 非有限； +// 2) ref 有限时 got 必须有限，再用 |got-ref| <= atol + rtol*|ref|； +// 3) 普通用例出现非预期 NaN/Inf（哪怕两侧相同）一律失败。 +// +// 模式（解析在 backend 注册与设备初始化之前完成）： +// （默认） 设备对拍套件；有失败 → 退出 1 +// --comparator-selftest 纯 host，逐个断言比较器判断，全部符合才退出 0 +// --inject-finite 设备对拍并注入有限误差，预期整轮失败（退出非 0，由外层脚本检查） +// --inject-nan 设备对拍并注入非预期 NaN，预期整轮失败 +// --probe-case NAME 设备上的探针用例（只分配+拷贝，不发 kernel），配合 MLU_DEBUG_CHECK=1 由脚本断言日志 +// --device N 选择设备 + +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include + +#include "infini_train/include/core/runtime/device_guard.h" +#include "infini_train/include/datatype.h" +#include "infini_train/include/device.h" +#include "infini_train/include/dispatcher.h" +#include "infini_train/include/tensor.h" + +#include "infini_train_mlu/backend.h" +#include "runtime/mlu_debug.h" + +namespace { + +using infini_train::DataType; +using infini_train::Device; +using infini_train::Tensor; + +constexpr cnrtRet_t CNRT_RET_SUCCESS_LOCAL = cnrtSuccess; +constexpr auto kCPU = Device::DeviceType::kCPU; +constexpr auto kMLU = Device::DeviceType::kPrivateUse1; + +bool g_failed = false; +bool g_inject_finite = false; +bool g_inject_nan = false; +int g_device = 0; + +Device Cpu() { return Device(kCPU, 0); } +Device Mlu() { return Device(kMLU, g_device); } + +std::shared_ptr MakeTensor(const std::vector &dims, const std::vector &host, Device dev) { + auto t = std::make_shared(dims, DataType::kFLOAT32, dev); + if (dev.type() == kCPU) { + std::memcpy(t->DataPtr(), host.data(), host.size() * sizeof(float)); + return t; + } + infini_train::core::GetDeviceGuardImpl(dev.type()) + ->Memcpy(t->DataPtr(), host.data(), host.size() * sizeof(float), infini_train::core::MemcpyKind::kH2D); + return t; +} + +std::shared_ptr MakeTensorI64(const std::vector &dims, const std::vector &host, Device dev) { + auto t = std::make_shared(dims, DataType::kINT64, dev); + if (dev.type() == kCPU) { + std::memcpy(t->DataPtr(), host.data(), host.size() * sizeof(int64_t)); + return t; + } + infini_train::core::GetDeviceGuardImpl(dev.type()) + ->Memcpy(t->DataPtr(), host.data(), host.size() * sizeof(int64_t), infini_train::core::MemcpyKind::kH2D); + return t; +} + +std::vector ToHost(const std::shared_ptr &t) { + std::vector host(t->NumElements()); + if (t->GetDevice().type() == kCPU) { + std::memcpy(host.data(), t->DataPtr(), host.size() * sizeof(float)); + return host; + } + infini_train::core::GetDeviceGuardImpl(t->GetDevice().type()) + ->Memcpy(host.data(), t->DataPtr(), host.size() * sizeof(float), infini_train::core::MemcpyKind::kD2H); + return host; +} + +std::vector RandVec(size_t n, float scale, uint32_t seed) { + std::mt19937 gen(seed); + std::normal_distribution dis(0.0f, scale); + std::vector v(n); + for (auto &x : v) { x = dis(gen); } + return v; +} + +// --------------------------------------------------------------------------- +// 比较器（R01）：非有限优先分类，只有两侧都有限时才用 atol+rtol +// --------------------------------------------------------------------------- +struct Verdict { + size_t bad = 0; + size_t worst = 0; + double max_abs = 0.0; + double max_rel = 0.0; + bool empty() const { return bad == 0; } +}; + +Verdict Evaluate(const std::vector &ref, const std::vector &got, double atol, double rtol, + bool allow_special_ref) { + Verdict v; + if (ref.size() != got.size()) { + v.bad = 1; + return v; + } + for (size_t i = 0; i < ref.size(); ++i) { + const double a = ref[i], b = got[i]; + const bool a_nan = std::isnan(a), b_nan = std::isnan(b); + const bool a_inf = std::isinf(a), b_inf = std::isinf(b); + + if (a_nan || a_inf) { + // 参考本身非有限：只有声明允许时才可接受，且必须同类同号 + const bool same = (a_nan && b_nan) || (a_inf && b_inf && ((a > 0) == (b > 0))); + if (!allow_special_ref || !same) { + ++v.bad; + v.worst = i; + v.max_abs = std::numeric_limits::infinity(); + } + continue; + } + if (b_nan || b_inf) { // 参考有限但结果非有限：一律失败 + ++v.bad; + v.worst = i; + v.max_abs = std::numeric_limits::infinity(); + continue; + } + const double d = std::fabs(a - b); + const double tol = atol + rtol * std::fabs(a); + const double r = d / (std::fabs(a) > 1e-12 ? std::fabs(a) : 1e-12); + if (d > tol) { + ++v.bad; + } + if (d > v.max_abs) { + v.max_abs = d; + v.worst = i; + } + if (r > v.max_rel) { + v.max_rel = r; + } + } + return v; +} + +bool Compare(const std::string &name, const std::vector &ref, const std::vector &got, double atol, + double rtol, bool allow_special_ref = false) { + const Verdict v = Evaluate(ref, got, atol, rtol, allow_special_ref); + const bool ok = v.empty(); + if (!ok) { + g_failed = true; + } + std::printf("%-42s %-5s n=%-6zu max_abs=%-11.3e max_rel=%-11.3e bad=%-3zu worst[%zu] ref=%.7g got=%.7g\n", + name.c_str(), ok ? "PASS" : "FAIL", ref.size(), v.max_abs, v.max_rel, v.bad, v.worst, + ref.empty() ? 0.0f : ref[v.worst], got.empty() ? 0.0f : got[v.worst]); + return ok; +} + +// 用解析期望值直接断言（不依赖 CPU 实现，R06） +bool CompareAnalytic(const std::string &name, const std::vector &expect, const std::vector &got, + double atol, double rtol) { + std::vector ref(expect.size()); + for (size_t i = 0; i < expect.size(); ++i) { ref[i] = static_cast(expect[i]); } + return Compare(name, ref, got, atol, rtol, /*allow_special_ref=*/true); +} + +// --------------------------------------------------------------------------- +// 纯 host：比较器自检（R02/R06），不初始化设备 +// --------------------------------------------------------------------------- +int RunComparatorSelftest() { + struct Case { + const char *name; + std::vector ref; + std::vector got; + bool allow_special; + bool expect_pass; + }; + const float inf = std::numeric_limits::infinity(); + const float nan_v = std::numeric_limits::quiet_NaN(); + const std::vector cases = { + {"finite equal", {1.0f}, {1.0f}, false, true}, + {"finite within tol", {1.0f}, {1.0f + 5e-5f}, false, true}, + {"finite beyond tol", {1.0f}, {1.01f}, false, false}, + {"ref +Inf got 0", {inf}, {0.0f}, false, false}, + {"ref -Inf got 0", {-inf}, {0.0f}, false, false}, + {"ref +Inf got -Inf", {inf}, {-inf}, false, false}, + {"ref +Inf got +Inf allowed", {inf}, {inf}, true, true}, + {"ref +Inf got +Inf not allowed", {inf}, {inf}, false, false}, + {"ref NaN got NaN not allowed", {nan_v}, {nan_v}, false, false}, + {"ref NaN got NaN allowed", {nan_v}, {nan_v}, true, true}, + {"ref NaN got 1 allowed", {nan_v}, {1.0f}, true, false}, + {"ref finite got NaN", {1.0f}, {nan_v}, true, false}, + {"ref finite got Inf", {1.0f}, {inf}, true, false}, + {"shape mismatch", {1.0f, 2.0f}, {1.0f}, false, false}, + {"tolerance boundary just inside", {2.0f}, {2.0f + 1e-5f + 1e-4f * 2.0f * 0.99f}, false, true}, + {"tolerance boundary just outside", {2.0f}, {2.0f + 1e-5f + 1e-4f * 2.0f * 1.5f}, false, false}, + }; + int failures = 0; + for (const auto &c : cases) { + const Verdict v = Evaluate(c.ref, c.got, 1e-5, 1e-4, c.allow_special); + const bool pass = v.empty(); + const bool ok = (pass == c.expect_pass); + if (!ok) { + ++failures; + } + std::printf("%-38s expect=%-4s actual=%-4s %s\n", c.name, c.expect_pass ? "PASS" : "FAIL", + pass ? "PASS" : "FAIL", ok ? "ok" : "MISMATCH"); + } + std::printf("comparator-selftest: %s (%d mismatch)\n", failures == 0 ? "PASS" : "FAIL", failures); + return failures == 0 ? 0 : 1; +} + +// --------------------------------------------------------------------------- +// 设备侧小算子用例 +// --------------------------------------------------------------------------- +void ApplyInjection(std::vector &mlu) { + if (mlu.empty()) { + return; + } + if (g_inject_finite) { + mlu[0] += 1.0f; + } + if (g_inject_nan) { + mlu[0] = std::numeric_limits::quiet_NaN(); + } +} + +void TestLinearForward(int64_t rows, int64_t in_f, int64_t out_f, bool transpose, bool with_bias) { + auto x = RandVec(static_cast(rows * in_f), 1.0f, 101); + auto w = RandVec(static_cast(in_f * out_f), 0.05f, 102); + auto b = RandVec(static_cast(out_f), 0.02f, 103); + const std::vector x_dims{rows, in_f}; + const std::vector w_dims + = transpose ? std::vector{out_f, in_f} : std::vector{in_f, out_f}; + auto run = [&](Device dev) { + auto xt = MakeTensor(x_dims, x, dev); + auto wt = MakeTensor(w_dims, w, dev); + auto bt = with_bias ? MakeTensor({out_f}, b, dev) : nullptr; + auto out = infini_train::Dispatcher::Instance().Call>({dev.type(), "LinearForward"}, xt, + wt, transpose, bt); + return ToHost(out); + }; + char name[128]; + std::snprintf(name, sizeof(name), "Linear fwd rows=%lld in=%lld out=%lld T=%d bias=%d", + static_cast(rows), static_cast(in_f), static_cast(out_f), + static_cast(transpose), static_cast(with_bias)); + auto cpu = run(Cpu()); + auto mlu = run(Mlu()); + ApplyInjection(mlu); + Compare(name, cpu, mlu, 1e-5, 1e-4); +} + +void TestLinearBackward(int64_t rows, int64_t in_f, int64_t out_f, bool transpose) { + auto x = RandVec(static_cast(rows * in_f), 1.0f, 201); + auto w = RandVec(static_cast(in_f * out_f), 0.05f, 202); + auto g = RandVec(static_cast(rows * out_f), 0.5f, 203); + const std::vector x_dims{rows, in_f}; + const std::vector w_dims + = transpose ? std::vector{out_f, in_f} : std::vector{in_f, out_f}; + const std::vector g_dims{rows, out_f}; + auto run = [&](Device dev) { + auto wt = MakeTensor(w_dims, w, dev); + auto gt = MakeTensor(g_dims, g, dev); + auto xt = MakeTensor(x_dims, x, dev); + auto dx = infini_train::Dispatcher::Instance().Call>( + {dev.type(), "LinearBackwardInput"}, wt, gt, transpose, in_f, out_f, x_dims); + auto dw = infini_train::Dispatcher::Instance().Call>( + {dev.type(), "LinearBackwardWeight"}, xt, gt, transpose, in_f, out_f); + auto db = infini_train::Dispatcher::Instance().Call>({dev.type(), "LinearBackwardBias"}, + gt, out_f); + return std::make_tuple(ToHost(dx), ToHost(dw), ToHost(db)); + }; + auto [cpu_dx, cpu_dw, cpu_db] = run(Cpu()); + auto [mlu_dx, mlu_dw, mlu_db] = run(Mlu()); + ApplyInjection(mlu_dx); + char tag[96]; + std::snprintf(tag, sizeof(tag), "rows=%lld in=%lld out=%lld T=%d", static_cast(rows), + static_cast(in_f), static_cast(out_f), static_cast(transpose)); + Compare(std::string("Linear bwd dx ") + tag, cpu_dx, mlu_dx, 1e-5, 1e-4); + Compare(std::string("Linear bwd dw ") + tag, cpu_dw, mlu_dw, 1e-5, 1e-4); + Compare(std::string("Linear bwd db ") + tag, cpu_db, mlu_db, 1e-5, 1e-4); +} + +// Pow：解析期望（R06）——0^0.5=0, 0^0=1, (-2)^9=-512, (-2)^0.5=NaN, 2^0.5=sqrt2, +// (-2)^3=-8, 0.5^-1=2, base=3/exp=2 => 9 +void TestPowAnalytic() { + struct Case { + float input; + float exponent; + bool base_is_scalar; + double expect; + const char *desc; + }; + const Case cases[] = { + {0.0f, 0.5f, false, 0.0, "0^0.5"}, + {0.0f, 0.0f, false, 1.0, "0^0"}, + {-2.0f, 9.0f, false, -512.0, "(-2)^9"}, + {-2.0f, 0.5f, false, std::numeric_limits::quiet_NaN(), "(-2)^0.5"}, + {2.0f, 0.5f, false, std::sqrt(2.0), "2^0.5"}, + {-2.0f, 3.0f, false, -8.0, "(-2)^3"}, + {0.5f, -1.0f, false, 2.0, "0.5^-1"}, + {2.0f, 3.0f, true, 9.0, "base=3 exp=2 -> 9"}, // tensor=2 是指数，scalar=3 是底数 + }; + for (const auto &c : cases) { + auto run = [&](Device dev) { + auto xt = MakeTensor({1}, {c.input}, dev); + auto out = infini_train::Dispatcher::Instance().Call>( + {dev.type(), "PowForward"}, xt, c.exponent, c.base_is_scalar); + return ToHost(out); + }; + auto cpu = run(Cpu()); + auto got = run(Mlu()); + char name[96]; + std::snprintf(name, sizeof(name), "Pow %s", c.desc); + // 解析期望与 CPU 参考双保险：先按解析值断言，再与 CPU 对拍 + CompareAnalytic(std::string(name) + " [analytic]", {c.expect}, got, 1e-6, 1e-5); + Compare(std::string(name) + " [vs cpu]", cpu, got, 1e-6, 1e-5, /*allow_special_ref=*/true); + } + + // scalar_is_base 反向：y=base^x, dy/dx = base^x*ln(base)；base=3, x=2, grad=2 -> 2*9*ln3 + { + const float base = 3.0f, xval = 2.0f, grad = 2.0f; + const double expect = static_cast(grad) * std::pow(base, xval) * std::log(base); + auto run = [&](Device dev) { + auto gt = MakeTensor({1}, {grad}, dev); + auto xt = MakeTensor({1}, {xval}, dev); + auto out = infini_train::Dispatcher::Instance().Call>({dev.type(), "PowBackward"}, + gt, xt, base, true); + return ToHost(out); + }; + CompareAnalytic("Pow bwd scalar_is_base=3 x=2 grad=2", {expect}, run(Mlu()), 1e-5, 1e-5); + } + + // Pow(x,3) 前向/反向，输入显式含 0 与负值 + { + const std::vector xs = {0.0f, -2.0f, 0.5f, 3.0f}; + const std::vector gs = {1.0f, 2.0f, -1.0f, 0.5f}; + std::vector exp_f(4), exp_b(4); + for (size_t i = 0; i < xs.size(); ++i) { + const double x = xs[i]; + exp_f[i] = x * x * x; + exp_b[i] = static_cast(gs[i]) * 3.0 * x * x; // x=0 处为 0 + } + auto run_fwd = [&](Device dev) { + auto xt = MakeTensor({4}, xs, dev); + auto out = infini_train::Dispatcher::Instance().Call>({dev.type(), "PowForward"}, + xt, 3.0f, false); + return ToHost(out); + }; + auto run_bwd = [&](Device dev) { + auto gt = MakeTensor({4}, gs, dev); + auto xt = MakeTensor({4}, xs, dev); + auto out = infini_train::Dispatcher::Instance().Call>({dev.type(), "PowBackward"}, + gt, xt, 3.0f, false); + return ToHost(out); + }; + CompareAnalytic("Pow(x,3) fwd incl x=0", exp_f, run_fwd(Mlu()), 1e-6, 1e-5); + CompareAnalytic("Pow(x,3) bwd incl x=0", exp_b, run_bwd(Mlu()), 1e-6, 1e-5); + } +} + +void TestCrossEntropy(int64_t bs, int64_t classes) { + auto logits = RandVec(static_cast(bs * classes), 1.0f, 401); + std::vector target(static_cast(bs)); + for (int64_t i = 0; i < bs; ++i) { target[i] = i % classes; } + auto run = [&](Device dev) { + auto lt = MakeTensor({bs, classes}, logits, dev); + auto tt = MakeTensorI64({bs}, target, dev); + auto out = infini_train::Dispatcher::Instance().Call>( + {dev.type(), "CrossEntropyForward"}, lt, tt); + return ToHost(out); + }; + Compare("CrossEntropy fwd (bs=4,cls=8)", run(Cpu()), run(Mlu()), 1e-5, 1e-4); +} + +void TestLayerNorm(int64_t rows, int64_t dim) { + auto x = RandVec(static_cast(rows * dim), 1.0f, 11); + auto w = RandVec(static_cast(dim), 0.1f, 12); + auto b = RandVec(static_cast(dim), 0.1f, 13); + auto run = [&](Device dev) { + auto xt = MakeTensor({rows, 1, dim}, x, dev); + auto wt = MakeTensor({dim}, w, dev); + auto bt = MakeTensor({dim}, b, dev); + auto res = infini_train::Dispatcher::Instance() + .Call, std::shared_ptr, std::shared_ptr>>( + {dev.type(), "LayerNormForward"}, xt, wt, bt, 1e-5f); + return ToHost(std::get<0>(res)); + }; + Compare("LayerNorm", run(Cpu()), run(Mlu()), 1e-5, 1e-4); +} + +void TestSoftmax(int64_t rows, int64_t cols) { + auto x = RandVec(static_cast(rows * cols), 1.0f, 21); + auto run = [&](Device dev) { + auto xt = MakeTensor({rows, cols}, x, dev); + auto out = infini_train::Dispatcher::Instance().Call>({dev.type(), "SoftmaxForward"}, + xt, int64_t{-1}); + return ToHost(out); + }; + Compare("Softmax(dim=-1)", run(Cpu()), run(Mlu()), 1e-5, 1e-4); +} + +void TestAddBroadcast(int64_t rows, int64_t cols) { + auto a = RandVec(static_cast(rows * cols), 1.0f, 31); + auto b = RandVec(static_cast(cols), 0.5f, 32); + auto run = [&](Device dev) { + auto at = MakeTensor({rows, cols}, a, dev); + auto bt = MakeTensor({cols}, b, dev); + auto out + = infini_train::Dispatcher::Instance().Call>({dev.type(), "AddForward"}, at, bt); + return ToHost(out); + }; + Compare("Add(broadcast)", run(Cpu()), run(Mlu()), 1e-5, 1e-4); +} + +void TestMulScalar(int64_t n, float s) { + auto x = RandVec(static_cast(n), 1.0f, 41); + auto run = [&](Device dev) { + auto xt = MakeTensor({n}, x, dev); + auto out = infini_train::Dispatcher::Instance().Call>({dev.type(), "MulScalarForward"}, + xt, s); + return ToHost(out); + }; + auto cpu = run(Cpu()); + auto mlu = run(Mlu()); + ApplyInjection(mlu); + Compare("MulScalar", cpu, mlu, 1e-5, 1e-4); +} + +void TestTranspose(int64_t b, int64_t h, int64_t t, int64_t d) { + auto x = RandVec(static_cast(b * h * t * d), 1.0f, 51); + auto run = [&](Device dev) { + auto xt = MakeTensor({b, h, t, d}, x, dev); + auto out = infini_train::Dispatcher::Instance().Call>({dev.type(), "TransposeForward"}, + xt, int64_t{1}, int64_t{2}); + return ToHost(out); + }; + Compare("Transpose(1,2)", run(Cpu()), run(Mlu()), 1e-5, 1e-4); +} + +// --------------------------------------------------------------------------- +// 探针用例（R05）：只分配 + 拷贝，不发 kernel；由脚本配合 MLU_DEBUG_CHECK=1 断言日志 +// --------------------------------------------------------------------------- +template +std::shared_ptr MakeRawTensor(const std::vector &dims, const std::vector &bits, DataType dt, + Device dev) { + auto t = std::make_shared(dims, dt, dev); + infini_train::core::GetDeviceGuardImpl(dev.type()) + ->Memcpy(t->DataPtr(), bits.data(), bits.size() * sizeof(Bits), infini_train::core::MemcpyKind::kH2D); + return t; +} + +int RunProbeCase(const std::string &name) { + using infini_train::core::mlu::DebugCheck; + if (name == "fp32_special") { + // [1, NaN, +Inf, -Inf, 0] -> non_finite=3, 有限 min=0 max=1 + auto t = MakeTensor({5}, + {1.0f, std::numeric_limits::quiet_NaN(), std::numeric_limits::infinity(), + -std::numeric_limits::infinity(), 0.0f}, + Mlu()); + DebugCheck(t, "probe_fp32_special"); + } else if (name == "first_nan") { + // 首元素 NaN：[NaN, 1, 2] -> non_finite=1, min=1 max=2（验证首个有限值初始化） + auto t = MakeTensor({3}, {std::numeric_limits::quiet_NaN(), 1.0f, 2.0f}, Mlu()); + DebugCheck(t, "probe_first_nan"); + } else if (name == "all_nonfinite") { + auto t = MakeTensor({3}, + {std::numeric_limits::quiet_NaN(), std::numeric_limits::infinity(), + -std::numeric_limits::infinity()}, + Mlu()); + DebugCheck(t, "probe_all_nonfinite"); + } else if (name == "bf16_bits") { + // 0x7c00 是合法有限 bf16（指数 0xf8），0x7f80 才是 +Inf，0x3f80=1.0，0x8000=-0 + std::vector bits = {0x7c00, 0x7f80, 0x3f80, 0x8000}; + auto t = MakeRawTensor({4}, bits, DataType::kBFLOAT16, Mlu()); + DebugCheck(t, "probe_bf16_bits"); + } else if (name == "fp16_bits") { + // 0x7c00 是 +Inf（fp16 指数 0x1f），0x3c00=1.0，0xfc00=-Inf + std::vector bits = {0x7c00, 0x3c00, 0xfc00}; + auto t = MakeRawTensor({3}, bits, DataType::kFLOAT16, Mlu()); + DebugCheck(t, "probe_fp16_bits"); + } else if (name == "int_wide") { + // 宽整数范围（R04）：UINT64_MAX 与 INT64_MIN/MAX 不得显示成 -1 或溢出 + std::vector u = {0, 18446744073709551615ULL}; + auto tu = MakeRawTensor({2}, u, DataType::kUINT64, Mlu()); + DebugCheck(tu, "probe_uint64"); + std::vector s = {INT64_MIN, INT64_MAX}; + auto ts = MakeRawTensor({2}, s, DataType::kINT64, Mlu()); + DebugCheck(ts, "probe_int64"); + } else { + std::printf("unknown probe case: %s\n", name.c_str()); + return 2; + } + return 0; +} + +} // namespace + +int main(int argc, char **argv) { + // 模式解析必须在 backend 注册/设备初始化之前（001-P2 v2 要求） + std::string probe_case; + std::string status_test; + bool comparator_selftest = false; + for (int i = 1; i < argc; ++i) { + const std::string arg(argv[i]); + if (arg == "--comparator-selftest") { + comparator_selftest = true; + } else if (arg == "--inject-finite") { + g_inject_finite = true; + } else if (arg == "--inject-nan") { + g_inject_nan = true; + } else if (arg == "--debug-status-test" && i + 1 < argc) { + status_test = argv[++i]; + } else if (arg == "--probe-case" && i + 1 < argc) { + probe_case = argv[++i]; + } else if (arg == "--device" && i + 1 < argc) { + g_device = std::atoi(argv[++i]); + } else { + std::printf("unknown arg: %s\n", argv[i]); + return 2; + } + } + + if (comparator_selftest) { + return RunComparatorSelftest(); + } + + // S03：受控失败路径测试。只调用既有 CheckSyncResult 的伪返回码，不把任何 queue 交给 CNRT， + // 也不触发真实设备故障。父脚本按"唯一标识 + 预期终止 + 无后续标记"判定。 + if (!status_test.empty()) { + // S03: controlled failure-path test. Calls the existing CheckSyncResult with a fabricated + // return code only; no queue is handed to CNRT and no real device fault is provoked. + if (status_test == "success") { + std::cout << "STATUS_TEST_BEGIN success" << std::endl; + infini_train::core::mlu::CheckSyncResult(cnrtSuccess, "debug-status-test-success"); + std::cout << "STATUS_TEST_END success (continued as expected)" << std::endl; + return 0; + } + if (status_test == "failure") { + std::cout << "STATUS_TEST_BEGIN failure" << std::endl; + std::cout << "STATUS_TEST_ABORT_MARKER fake-cnrt-failure-632015" << std::endl; + std::fflush(stdout); + infini_train::core::mlu::CheckSyncResult(static_cast(632015), "debug-status-test-failure"); + std::cout << "STATUS_TEST_AFTER_FATAL should-not-appear" << std::endl; + return 0; + } + std::cout << "unknown --debug-status-test value: " << status_test << std::endl; + return 2; + } + + infini_train::backends::mlu::RegisterBackend(); + std::printf("op_test: device=%d inject_finite=%d inject_nan=%d\n", g_device, static_cast(g_inject_finite), + static_cast(g_inject_nan)); + + if (!probe_case.empty()) { + return RunProbeCase(probe_case); + } + + for (bool transpose : {true, false}) { + TestLinearForward(3, 5, 7, transpose, true); + TestLinearForward(3, 5, 7, transpose, false); + TestLinearBackward(3, 5, 7, transpose); + } + TestPowAnalytic(); + TestCrossEntropy(4, 8); + TestLayerNorm(8, 16); + TestSoftmax(8, 16); + TestAddBroadcast(8, 16); + TestMulScalar(64, 0.5f); + TestTranspose(1, 2, 4, 8); + + std::printf("RESULT: %s\n", g_failed ? "FAIL" : "PASS"); + return g_failed ? 1 : 0; // 注入模式下同样如实返回非 0，由外层脚本检查预期退出码 +} diff --git a/backends/mlu/third_party/bang_host/README.md b/backends/mlu/third_party/bang_host/README.md new file mode 100644 index 0000000..263c880 --- /dev/null +++ b/backends/mlu/third_party/bang_host/README.md @@ -0,0 +1,38 @@ +# bang_host 头文件（vendor 拷贝） + +`bang_fp16.h` / `bang_bf16.h` 拷自 Neuware 安装： + + /usr/local/neuware/lib/clang/11.1.0/include/ (Neuware CNRT 7.7.3 / cncc 5.7.6) + +## 为什么拷贝 + +host 侧（g++）编译单元需要 `__half` / `bfloat16_t` 类型定义 +（`mlu/runtime/mlu_dispatch.h`），但这两个头只在 cncc 的 clang 资源目录里， +该目录同时含有 clang 专用的 `stddef.h` 等内建头。若把整个目录加进 g++ 的 +include 路径，clang 的 `stddef.h` 会遮蔽 GCC 自带版本（`__has_feature` 是 +clang 专有扩展，GCC 预处理器直接报错），导致所有 host 编译单元失败。 + +这两个头本身是自包含的（`bang_bf16.h` 仅 include `bang_fp16.h`，后者无 +include），host 路径（`__GNUC__` 分支）不依赖任何 clang 内建头，因此拷贝 +出来单独使用是安全的。 + +`.mlu` 设备代码不受影响：cncc 自动从自己的资源目录找到 `` 等头， +无需 `-I`。 + +## 许可证 + +两头部为 Cambricon 的 permissive license（见文件头注释，允许拷贝/修改/ +再分发，需保留版权声明）。 + +## 本地修改（与原版的唯一差异） + +host 分支（非 `__BANGCC__`）的 `__BANG_HOSTDEVICE__` / +`__BANG_HOSTDEVICE_BF16__` 宏从空改为 `inline`：原版假设 host 侧只有 +一个编译单元包含这些头，多 TU 包含时全部 free operator / 成员函数定义 +会 multiple definition 链接错误；设备（cncc）分支本就带 `inline`，此改动 +只是让 host 分支与设备分支一致。 + +## 更新时机 + +Neuware 大版本升级后，对比原目录中的同名文件，有变化则同步更新。 +host 侧只用到位布局稳定的 16 位浮点类型定义，版本漂移风险很低。 diff --git a/backends/mlu/third_party/bang_host/bang_bf16.h b/backends/mlu/third_party/bang_host/bang_bf16.h new file mode 100644 index 0000000..ffbf94b --- /dev/null +++ b/backends/mlu/third_party/bang_host/bang_bf16.h @@ -0,0 +1,829 @@ +/*===-----------------------------------------------------------------------=== + * + * Copyright (C) [2022] by Cambricon, Inc. + * + * Permission is hereby granted, free of charge, to any person obtaining a copy + * of this software and associated documentation files (the "Software"), to deal + * in the Software without restriction, including without limitation the rights + * to use, copy, modify, merge, publish, distribute, sublicense, and/or sell + * copies of the Software, and to permit persons to whom the Software is + * furnished to do so, subject to the following conditions: + * + * The above copyright notice and this permission notice shall be included in + * all copies or substantial portions of the Software. + * + * THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR + * IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, + * FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE + * AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER + * LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, + * OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN + * THE SOFTWARE. + * + *===-----------------------------------------------------------------------=== + */ +#if !defined(__BANG_BF16_H__) +#define __BANG_BF16_H__ + +#include "bang_fp16.h" + +#if defined(__BANGCC__) +#define __BANG_HOSTDEVICE_BF16__ __mlu_host__ __mlu_func__ +#define __BANG_HOSTDEVICE_BF16_DECL__ static __mlu_host__ __mlu_func__ +#else /* !defined(__BANGCC__) */ +#if defined(__GNUC__) +#define __BANG_HOSTDEVICE_BF16_DECL__ static __attribute__((unused)) +#else +#define __BANG_HOSTDEVICE_BF16_DECL__ static +#endif /* defined(__GNUC__) */ +#define __BANG_HOSTDEVICE_BF16__ inline +#endif /* defined(__BANGCC_) */ + +#if !defined(_MSC_VER) && __cplusplus >= 201103L +# define __CPP_VERSION_AT_LEAST_11_BF16 +#elif _MSC_FULL_VER >= 190024210 && _MSVC_LANG >= 201103L +# define __CPP_VERSION_AT_LEAST_11_BF16 +#endif + +/* Set up structure-alignment attribute */ +#if defined(__CPP_VERSION_AT_LEAST_11_BF16) +#define __BANG_ALIGN_BF16__(n) alignas(n) /* C++11 kindly gives us a keyword for this */ +#define __BANG_BF16_CONSTEXPR__ constexpr +#else /* defined(__CPP_VERSION_AT_LEAST_11_BF16)*/ +#if defined(__GNUC__) +#define __BANG_ALIGN_BF16__(n) __attribute__ ((aligned(n))) +#elif defined(_MSC_VER) +#define __BANG_ALIGN_BF16__(n) __declspec(align(n)) +#else +#define __BANG_ALIGN_BF16__(n) +#endif /* defined(__GNUC__) */ +#define __BANG_BF16_CONSTEXPR__ +#endif /* defined(__CPP_VERSION_AT_LEAST_11_BF16) */ + +/* Macros to allow bang_bfloat16 to be used by inline assembly */ +#define __BFLOAT16_TO_US(var) *(reinterpret_cast(&(var))) +#define __BFLOAT16_TO_CUS(var) *(reinterpret_cast(&(var))) +#define __FLOAT_TO_UINT(var) *(reinterpret_cast(&(var))) +#define __UINT_TO_FLOAT(var) *(reinterpret_cast(&(var))) + +typedef struct __BANG_ALIGN_BF16__(2) { + unsigned short x; +} __bang_bfloat16_raw; + +/* All other definitions in this file are only visible to C++ compilers */ +#if defined(__cplusplus) + +/* Hide GCC member initialization list warnings because of host/device in-function init requirement */ +#if defined(__GNUC__) +#if __GNUC__ > 4 || (__GNUC__ == 4 && __GNUC_MINOR__ >= 6) +#pragma GCC diagnostic push +#pragma GCC diagnostic ignored "-Wstrict-aliasing" +#pragma GCC diagnostic ignored "-Weffc++" +#endif /* __GNUC__ > 4 || (__GNUC__ == 4 && __GNUC_MINOR__ >= 6) */ +#endif /* defined(__GNUC__) */ + +/* class' : multiple assignment operators specified + The class has multiple assignment operators of a single type. This warning is informational */ +#if defined(_MSC_VER) && _MSC_VER >= 1500 +#pragma warning( push ) +#pragma warning( disable:4522 ) +#endif /* defined(__GNUC__) */ + +struct __BANG_ALIGN_BF16__(2) __bang_bfloat16; +__BANG_HOSTDEVICE_BF16_DECL__ __bang_bfloat16 __float2bfloat16__(const float a); +__BANG_HOSTDEVICE_BF16_DECL__ float __bfloat162float__(const __bang_bfloat16 a); +__BANG_HOSTDEVICE_BF16_DECL__ __bang_bfloat16 __short2bfloat16_rn__(const short int i); +__BANG_HOSTDEVICE_BF16_DECL__ __bang_bfloat16 __ushort2bfloat16_rn__(const unsigned short int i); +__BANG_HOSTDEVICE_BF16_DECL__ __bang_bfloat16 __int2bfloat16_rn__(const int i); +__BANG_HOSTDEVICE_BF16_DECL__ int __bfloat162int_rz__(const __bang_bfloat16 h); +__BANG_HOSTDEVICE_BF16_DECL__ short int __bfloat162short_rz__(const __bang_bfloat16 h); +__BANG_HOSTDEVICE_BF16_DECL__ unsigned int __bfloat162uint_rz__(const __bang_bfloat16 h); +__BANG_HOSTDEVICE_BF16_DECL__ __bang_bfloat16 __uint2bfloat16_rn__(const unsigned int i); +__BANG_HOSTDEVICE_BF16_DECL__ unsigned short int __bfloat162ushort_rz__(const __bang_bfloat16 h); + +__BANG_HOSTDEVICE_BF16_DECL__ __bang_bfloat16 __hneg(const __bang_bfloat16 a); +__BANG_HOSTDEVICE_BF16_DECL__ __bang_bfloat16 __hsub(const __bang_bfloat16 a, const __bang_bfloat16 b); +__BANG_HOSTDEVICE_BF16_DECL__ __bang_bfloat16 __hadd(const __bang_bfloat16 a, const __bang_bfloat16 b); +__BANG_HOSTDEVICE_BF16_DECL__ __bang_bfloat16 __hmul(const __bang_bfloat16 a, const __bang_bfloat16 b); +__BANG_HOSTDEVICE_BF16_DECL__ __bang_bfloat16 __hdiv(const __bang_bfloat16 a, const __bang_bfloat16 b); +__BANG_HOSTDEVICE_BF16_DECL__ bool __heq(const __bang_bfloat16 a, const __bang_bfloat16 b); +__BANG_HOSTDEVICE_BF16_DECL__ bool __hneu(const __bang_bfloat16 a, const __bang_bfloat16 b); +__BANG_HOSTDEVICE_BF16_DECL__ bool __hgt(const __bang_bfloat16 a, const __bang_bfloat16 b); +__BANG_HOSTDEVICE_BF16_DECL__ bool __hlt(const __bang_bfloat16 a, const __bang_bfloat16 b); +__BANG_HOSTDEVICE_BF16_DECL__ bool __hge(const __bang_bfloat16 a, const __bang_bfloat16 b); +__BANG_HOSTDEVICE_BF16_DECL__ bool __hle(const __bang_bfloat16 a, const __bang_bfloat16 b); + +__BANG_HOSTDEVICE_BF16_DECL__ __bang_bfloat16 __ll2bfloat16_rn(const long long int i); +__BANG_HOSTDEVICE_BF16_DECL__ long long int __bfloat162ll_rz(const __bang_bfloat16 h); +__BANG_HOSTDEVICE_BF16_DECL__ __bang_bfloat16 __ull2bfloat16_rn(const unsigned long long int i); +__BANG_HOSTDEVICE_BF16_DECL__ unsigned long long int __bfloat162ull_rz(const __bang_bfloat16 h); + +struct __BANG_ALIGN_BF16__(2) __bang_bfloat16 { +protected: + unsigned short __x; + +public: +#if defined(__CPP_VERSION_AT_LEAST_11_BF16) + __bang_bfloat16() = default; +#else + __BANG_HOSTDEVICE_BF16__ __bang_bfloat16() { } +#endif /* defined(__CPP_VERSION_AT_LEAST_11_BF16) */ + + /* Convert to/from __bang_bfloat16_raw */ + __BANG_HOSTDEVICE_BF16__ __BANG_BF16_CONSTEXPR__ __bang_bfloat16(const __bang_bfloat16_raw &hr) : __x(hr.x) { } + __BANG_HOSTDEVICE_BF16__ __bang_bfloat16 &operator=(const __bang_bfloat16_raw &hr) { __x = hr.x; return *this; } + __BANG_HOSTDEVICE_BF16__ volatile __bang_bfloat16 &operator=(const __bang_bfloat16_raw &hr) volatile { __x = hr.x; return *this; } + __BANG_HOSTDEVICE_BF16__ volatile __bang_bfloat16 &operator=(const volatile __bang_bfloat16_raw &hr) volatile { __x = hr.x; return *this; } + __BANG_HOSTDEVICE_BF16__ operator __bang_bfloat16_raw() const { __bang_bfloat16_raw ret; ret.x = __x; return ret; } + __BANG_HOSTDEVICE_BF16__ operator __bang_bfloat16_raw() const volatile { __bang_bfloat16_raw ret; ret.x = __x; return ret; } + + // /* Construct from float */ + __BANG_HOSTDEVICE_BF16__ __bang_bfloat16(const float f) { __x = __float2bfloat16__(f).__x; } + + __BANG_HOSTDEVICE_BF16__ operator float() const { return __bfloat162float__(*this); } + __BANG_HOSTDEVICE_BF16__ __bang_bfloat16 &operator=(const float f) { __x = __float2bfloat16__(f).__x; return *this; } + + // /* Allow automatic construction from types supported natively in hardware */ + // /* Note we do avoid constructor init-list because of special host/device compilation rules */ + __BANG_HOSTDEVICE_BF16__ __bang_bfloat16(short val) { __x = __short2bfloat16_rn__(val).__x; } + __BANG_HOSTDEVICE_BF16__ __bang_bfloat16(unsigned short val) { __x = __ushort2bfloat16_rn__(val).__x; } + __BANG_HOSTDEVICE_BF16__ __bang_bfloat16(int val) { __x = __int2bfloat16_rn__(val).__x; } + __BANG_HOSTDEVICE_BF16__ __bang_bfloat16(unsigned int val) { __x = __uint2bfloat16_rn__(val).__x; } + + __BANG_HOSTDEVICE_BF16__ __bang_bfloat16(const __half f) { +#if defined(__BANG_ARCH__) + asm("cvt.gpr.rn.bf16.f16 %0, %1;" : "+r"(__x) : "r"(__BFLOAT16_TO_CUS(f))); +#else + __x = __float2bfloat16__(__half2float(f)).__x; +#endif + } + + /* Allow automatic casts to supported builtin types, matching all that are permitted with float */ + __BANG_HOSTDEVICE_BF16__ operator short() const { return __bfloat162short_rz__(*this); } + __BANG_HOSTDEVICE_BF16__ __bang_bfloat16 &operator=(short val) { __x = __short2bfloat16_rn__(val).__x; return *this; } + + __BANG_HOSTDEVICE_BF16__ operator unsigned short() const { return __bfloat162ushort_rz__(*this); } + __BANG_HOSTDEVICE_BF16__ __bang_bfloat16 &operator=(unsigned short val) { __x = __ushort2bfloat16_rn__(val).__x; return *this; } + + __BANG_HOSTDEVICE_BF16__ operator int() const { return __bfloat162int_rz__(*this); } + __BANG_HOSTDEVICE_BF16__ __bang_bfloat16 &operator=(int val) { __x = __int2bfloat16_rn__(val).__x; return *this; } + + __BANG_HOSTDEVICE_BF16__ operator unsigned int() const { return __bfloat162uint_rz__(*this); } + __BANG_HOSTDEVICE_BF16__ __bang_bfloat16 &operator=(unsigned int val) { __x = __uint2bfloat16_rn__(val).__x; return *this; } + + /* Boolean conversion - note both 0 and -0 must return false */ + __BANG_HOSTDEVICE_BF16__ operator bool() const { return (__x & 0x7FFF) != 0; } + + __BANG_HOSTDEVICE_BF16__ __bang_bfloat16(const long val) { + if (sizeof(long) == sizeof(long long)) { + __x = __ll2bfloat16_rn(static_cast(val)).__x; + } else { + __x = __int2bfloat16_rn__(static_cast(val)).__x; + } + } + + __BANG_HOSTDEVICE_BF16__ __bang_bfloat16(const unsigned long val) { + if (sizeof(unsigned long) == sizeof(unsigned long long)) { + __x = __ull2bfloat16_rn(static_cast(val)).__x; + } else { + __x = __uint2bfloat16_rn__(static_cast(val)).__x; + } + } + + __BANG_HOSTDEVICE_BF16__ __bang_bfloat16(long long val) { __x = __ll2bfloat16_rn(val).__x; } + __BANG_HOSTDEVICE_BF16__ __bang_bfloat16(unsigned long long val) { __x = __ull2bfloat16_rn(val).__x; } + + __BANG_HOSTDEVICE_BF16__ operator long() const { + long retval; + if (sizeof(long) == sizeof(long long)) { + retval = static_cast(__bfloat162ll_rz(*this)); + } else { + retval = static_cast(__bfloat162int_rz__(*this)); + } + return retval; + } + + __BANG_HOSTDEVICE_BF16__ operator unsigned long() const { + unsigned long retval; + if (sizeof(unsigned long) == sizeof(unsigned long long)) { + retval = static_cast(__bfloat162ull_rz(*this)); + } else { + retval = static_cast(__bfloat162uint_rz__(*this)); + } + return retval; + } + + __BANG_HOSTDEVICE_BF16__ operator long long() const { return __bfloat162ll_rz(*this); } + __BANG_HOSTDEVICE_BF16__ operator unsigned long long() const { return __bfloat162ull_rz(*this); } + + __BANG_HOSTDEVICE_BF16__ __bang_bfloat16 &operator=(long long val) { __x = __ll2bfloat16_rn(val).__x; return *this; } + __BANG_HOSTDEVICE_BF16__ __bang_bfloat16 &operator=(long val) { __x = __bang_bfloat16(val).__x; return *this; } + + __BANG_HOSTDEVICE_BF16__ __bang_bfloat16 &operator=(unsigned long long val) { __x = __ull2bfloat16_rn(val).__x; return *this; } + __BANG_HOSTDEVICE_BF16__ __bang_bfloat16 &operator=(unsigned long val) { __x = __bang_bfloat16(val).__x; return *this; } + + __BANG_HOSTDEVICE_BF16__ __bang_bfloat16 &operator=(__half f) { __x = __bang_bfloat16(f).__x; return *this; } +}; + +/* Some basic arithmetic operations expected of a builtin */ +__BANG_HOSTDEVICE_BF16_DECL__ __bang_bfloat16 operator+(const __bang_bfloat16 &lh, const __bang_bfloat16 &rh) { return __hadd(lh, rh); } +__BANG_HOSTDEVICE_BF16_DECL__ __bang_bfloat16 operator-(const __bang_bfloat16 &lh, const __bang_bfloat16 &rh) { return __hsub(lh, rh); } +__BANG_HOSTDEVICE_BF16_DECL__ __bang_bfloat16 operator*(const __bang_bfloat16 &lh, const __bang_bfloat16 &rh) { return __hmul(lh, rh); } +__BANG_HOSTDEVICE_BF16_DECL__ __bang_bfloat16 operator/(const __bang_bfloat16 &lh, const __bang_bfloat16 &rh) { return __hdiv(lh, rh); } + +__BANG_HOSTDEVICE_BF16_DECL__ __bang_bfloat16 &operator+=(__bang_bfloat16 &lh, const __bang_bfloat16 &rh) { lh = __hadd(lh, rh); return lh; } +__BANG_HOSTDEVICE_BF16_DECL__ __bang_bfloat16 &operator-=(__bang_bfloat16 &lh, const __bang_bfloat16 &rh) { lh = __hsub(lh, rh); return lh; } +__BANG_HOSTDEVICE_BF16_DECL__ __bang_bfloat16 &operator*=(__bang_bfloat16 &lh, const __bang_bfloat16 &rh) { lh = __hmul(lh, rh); return lh; } +__BANG_HOSTDEVICE_BF16_DECL__ __bang_bfloat16 &operator/=(__bang_bfloat16 &lh, const __bang_bfloat16 &rh) { lh = __hdiv(lh, rh); return lh; } + +/* Note for increment and decrement we use the raw value 0x3F80 equating to bang_bfloat16(1.0f), to avoid the extra conversion */ +__BANG_HOSTDEVICE_BF16_DECL__ __bang_bfloat16 &operator++(__bang_bfloat16 &h) { __bang_bfloat16_raw one; one.x = 0x3F80; h += one; return h; } +__BANG_HOSTDEVICE_BF16_DECL__ __bang_bfloat16 &operator--(__bang_bfloat16 &h) { __bang_bfloat16_raw one; one.x = 0x3F80; h -= one; return h; } +__BANG_HOSTDEVICE_BF16_DECL__ __bang_bfloat16 operator++(__bang_bfloat16 &h, const int ignored) { + // ignored on purpose. Parameter only needed to distinguish the function declaration from other types of operators. + static_cast(ignored); + + const __bang_bfloat16 ret = h; + __bang_bfloat16_raw one; + one.x = 0x3F80; + h += one; + return ret; +} +__BANG_HOSTDEVICE_BF16_DECL__ __bang_bfloat16 operator--(__bang_bfloat16 &h, const int ignored) +{ + // ignored on purpose. Parameter only needed to distinguish the function declaration from other types of operators. + static_cast(ignored); + + const __bang_bfloat16 ret = h; + __bang_bfloat16_raw one; + one.x = 0x3F80; + h -= one; + return ret; +} +/* Unary plus and inverse operators */ +__BANG_HOSTDEVICE_BF16_DECL__ __bang_bfloat16 operator+(const __bang_bfloat16 &h) { return h; } +__BANG_HOSTDEVICE_BF16_DECL__ __bang_bfloat16 operator-(const __bang_bfloat16 &h) { return __hneg(h); } + +/* Some basic comparison operations to make it look like a builtin */ +__BANG_HOSTDEVICE_BF16_DECL__ bool operator==(const __bang_bfloat16 &lh, const __bang_bfloat16 &rh) { return __heq(lh, rh); } +__BANG_HOSTDEVICE_BF16_DECL__ bool operator!=(const __bang_bfloat16 &lh, const __bang_bfloat16 &rh) { return __hneu(lh, rh); } +__BANG_HOSTDEVICE_BF16_DECL__ bool operator> (const __bang_bfloat16 &lh, const __bang_bfloat16 &rh) { return __hgt(lh, rh); } +__BANG_HOSTDEVICE_BF16_DECL__ bool operator< (const __bang_bfloat16 &lh, const __bang_bfloat16 &rh) { return __hlt(lh, rh); } +__BANG_HOSTDEVICE_BF16_DECL__ bool operator>=(const __bang_bfloat16 &lh, const __bang_bfloat16 &rh) { return __hge(lh, rh); } +__BANG_HOSTDEVICE_BF16_DECL__ bool operator<=(const __bang_bfloat16 &lh, const __bang_bfloat16 &rh) { return __hle(lh, rh); } + +__BANG_HOSTDEVICE_BF16__ __half::__half(const __bang_bfloat16 f) { +#if defined(__BANG_ARCH__) + asm("cvt.gpr.rn.f16.bf16 %0, %1;" : "+r"(__x) : "r"(__BFLOAT16_TO_CUS(f))); +#else + __x = __float2half_rn__(__bfloat162float__(f)).__x; +#endif +} + +__BANG_HOSTDEVICE_BF16__ __half &__half::operator=(const __bang_bfloat16 f) { __x = __half(f).__x; return *this; } + +/* Restore warning for multiple assignment operators */ +#if defined(_MSC_VER) && _MSC_VER >= 1500 +#pragma warning( pop ) +#endif /* defined(_MSC_VER) && _MSC_VER >= 1500 */ + +/* Restore -Weffc++ warnings from here on */ +#if defined(__GNUC__) +#if __GNUC__ > 4 || (__GNUC__ == 4 && __GNUC_MINOR__ >= 6) +#pragma GCC diagnostic pop +#endif /* __GNUC__ > 4 || (__GNUC__ == 4 && __GNUC_MINOR__ >= 6) */ +#endif /* defined(__GNUC__) */ + +#undef __BANG_HOSTDEVICE_BF16__ +#undef __BANG_ALIGN_BF16__ +#undef __BANG_BF16_CONSTEXPR__ + +__BANG_HOSTDEVICE_BF16_DECL__ unsigned int __internal_float_as_uint(const float f) { + unsigned int u = *(reinterpret_cast(&(f))); + return u; +} + +__BANG_HOSTDEVICE_BF16_DECL__ float __internal_uint_as_float(unsigned int u) { + float f = *(reinterpret_cast(&u)); + return f; +} + +__BANG_HOSTDEVICE_BF16_DECL__ unsigned short __internal_float2bfloat16(const float f, unsigned int &sign, unsigned int &remainder) +{ + unsigned int x; + + x = __internal_float_as_uint(f); + + if ((x & 0x7fffffffU) > 0x7f800000U) { + sign = 0U; + remainder = 0U; + return static_cast(0x7fffU); + } + sign = x >> 31U; + remainder = x << 16U; + return static_cast(x >> 16U); +} + +__BANG_HOSTDEVICE_BF16_DECL__ __bang_bfloat16 __float2bfloat16__(const float a) +{ + __bang_bfloat16 val; +#if defined(__BANG_ARCH__) + __asm__("cvt.gpr.rn.bf16.f32 \t %[dst], %[src];\n\t" + : [dst] "+r"(__BFLOAT16_TO_US(val)) + : [src] "r"(a)); +#else + __bang_bfloat16_raw r; + unsigned int sign = 0U; + unsigned int remainder = 0U; + r.x = __internal_float2bfloat16(a, sign, remainder); + if ((remainder > 0x80000000U) || ((remainder == 0x80000000U) && ((r.x & 0x1U) != 0U))) { + r.x++; + } + val = r; +#endif + return val; +} +__BANG_HOSTDEVICE_BF16_DECL__ __bang_bfloat16 __float2bfloat16_rn__(const float a) +{ + __bang_bfloat16 val; +#if defined(__BANG_ARCH__) + __asm__("cvt.gpr.rn.bf16.f32 \t %[dst], %[src];\n\t" + : [dst] "+r"(__BFLOAT16_TO_US(val)) + : [src] "r"(a)); +#else + __bang_bfloat16_raw r; + unsigned int sign = 0U; + unsigned int remainder = 0U; + r.x = __internal_float2bfloat16(a, sign, remainder); + if ((remainder > 0x80000000U) || ((remainder == 0x80000000U) && ((r.x & 0x1U) != 0U))) { + r.x++; + } + val = r; +#endif + return val; +} + +__BANG_HOSTDEVICE_BF16_DECL__ float __internal_bfloat162float(const unsigned short h) +{ + float f; +#if defined(__BANG_ARCH__) + __asm__("cvt.gpr.rn.f32.bf16 \t %[dst], %[src];\n\t" + : [dst] "+r"(f) + : [src] "r"(h)); +#else + unsigned int u = static_cast(h) << 16; + f = __internal_uint_as_float(u); +#endif + return f; +} + +__BANG_HOSTDEVICE_BF16_DECL__ float __bfloat162float__(const __bang_bfloat16 a) +{ + return __internal_bfloat162float(static_cast<__bang_bfloat16_raw>(a).x); +} + +__BANG_HOSTDEVICE_BF16_DECL__ int __internal_bfloat162int_rz(const __bang_bfloat16 h) +{ + float f = __bfloat162float__(h); + int i; +#if defined(__BANG_ARCH__) + if (((*(int *)&f) & 0x7FFFFFFF) - 0x7F800000 > 0) { + f = 0; + } + + __asm__("cvt.gpr.tz.s32.f32 \t %[dst], %[src];\n\t" + : [dst] "+r"(i) + : [src] "r"(f)); +#else + const int max_val = (int)0x7fffffffU; + const int min_val = (int)0x80000000U; + const unsigned short bits = static_cast(static_cast<__bang_bfloat16_raw>(h).x << 1U); + // saturation fixup + if (bits > (unsigned short)0xFF00U) { + // NaN + i = 0; + } else if (f >= static_cast(max_val)) { + // saturate maximum + i = max_val; + } else if (f < static_cast(min_val)) { + // saturate minimum + i = min_val; + } else { + i = static_cast(f); + } +#endif + return i; +} + +__BANG_HOSTDEVICE_BF16_DECL__ int __bfloat162int_rz__(const __bang_bfloat16 h) +{ +#if defined(__BANG_ARCH__) + int val; + __asm__("cvt.gpr.tz.s32.bf16 \t %[dst], %[src];\n\t" + : [dst] "+r"(val) + : [src] "r"(__BFLOAT16_TO_CUS(h))); + return val; +#else + return __internal_bfloat162int_rz(h); +#endif +} + +__BANG_HOSTDEVICE_BF16_DECL__ __bang_bfloat16 __int2bfloat16_rn__(const int i) +{ +#if defined(__BANG_ARCH__) + __bang_bfloat16 val; + __asm__("cvt.gpr.rn.bf16.s32 \t %[dst], %[src];\n\t" + : [dst] "+r"(__BFLOAT16_TO_US(val)) + : [src] "r"(i)); + return val; +#else + const float d = static_cast(i); + return __float2bfloat16__(d); +#endif +} + +__BANG_HOSTDEVICE_BF16_DECL__ short int __bfloat162short_rz__(const __bang_bfloat16 h) +{ + short int val; +#if defined(__BANG_ARCH__) + __asm__("cvt.gpr.tz.s16.bf16 \t %[dst], %[src]; \n\t" + : [dst] "+r"(__BFLOAT16_TO_US(val)) + : [src] "r"(__BFLOAT16_TO_CUS(h))); +#else + const float f = __bfloat162float__(h); + const short int max_val = (short int)0x7fffU; + const short int min_val = (short int)0x8000U; + const unsigned short bits = static_cast(static_cast<__bang_bfloat16_raw>(h).x << 1U); + // saturation fixup + if (bits > (unsigned short)0xFF00U) { + // NaN + val = 0; + } else if (f > static_cast(max_val)) { + // saturate maximum + val = max_val; + } else if (f < static_cast(min_val)) { + // saturate minimum + val = min_val; + } else { + val = static_cast(f); + } +#endif + return val; +} + +__BANG_HOSTDEVICE_BF16_DECL__ __bang_bfloat16 __short2bfloat16_rn__(const short int i) +{ +#if defined(__BANG_ARCH__) + __bang_bfloat16 val; + __asm__("cvt.gpr.rn.bf16.s16 \t %[dst], %[src]; \n\t" + : [dst] "+r"(__BFLOAT16_TO_US(val)) + : [src] "r"(i)); + return val; +#else + const float f = static_cast(i); + return __float2bfloat16_rn__(f); +#endif +} + +__BANG_HOSTDEVICE_BF16_DECL__ unsigned int __internal_bfloat162uint_rz(const __bang_bfloat16 h) +{ + float f = __bfloat162float__(h); + unsigned int i; +#if defined(__BANG_ARCH__) + if (((*(int *)&f) & 0x7FFFFFFF) - 0x7F800000 > 0) { + f = 0; + } + + __asm__("cvt.gpr.tz.u32.f32 \t %[dst], %[src];\n\t" + : [dst] "+r"(i) + : [src] "r"(f)); +#else + const unsigned int max_val = 0xffffffffU; + const unsigned int min_val = 0U; + const unsigned short bits = static_cast(static_cast<__bang_bfloat16_raw>(h).x << 1U); + // saturation fixup + if (bits > (unsigned short)0xFF00U) { + // NaN + i = 0U; + } else if (f >= static_cast(max_val)) { + // saturate maximum + i = max_val; + } else if (f < static_cast(min_val)) { + // saturate minimum + i = min_val; + } else { + i = static_cast(f); + } +#endif + return i; +} + +__BANG_HOSTDEVICE_BF16_DECL__ unsigned int __bfloat162uint_rz__(const __bang_bfloat16 h) +{ + return __internal_bfloat162uint_rz(h); +} + +__BANG_HOSTDEVICE_BF16_DECL__ __bang_bfloat16 __uint2bfloat16_rn__(const unsigned int i) +{ +#if defined(__BANG_ARCH__) + __bang_bfloat16 val; + __asm__("cvt.gpr.rn.bf16.u32 \t %[dst], %[src];\n\t" + : [dst] "+r"(__BFLOAT16_TO_US(val)) + : [src] "r"(i)); + return val; +#else + const float d = static_cast(i); + return __float2bfloat16__(d); +#endif +} + +__BANG_HOSTDEVICE_BF16_DECL__ unsigned short int __bfloat162ushort_rz__(const __bang_bfloat16 h) +{ + unsigned short int val; +#if defined(__BANG_ARCH__) + float tmp; + __asm__("sll.gpr.u32 \t %[dst], %[src], 16; \n\t" + : [dst] "+r"(tmp) + : [src] "r"(__BFLOAT16_TO_CUS(h))); + + if (((*(int *)&tmp) & 0x7FFFFFFF) - 0x7F800000 > 0) { + tmp = 0; + } + + __asm__("cvt.gpr.tz.u16.f32 \t %[dst], %[src]; \n\t" + : [dst] "+r"(__BFLOAT16_TO_US(val)) + : [src] "r"(tmp)); +#else + const float f = __bfloat162float__(h); + const unsigned short int max_val = 0xffffU; + const unsigned short int min_val = 0U; + const unsigned short bits = static_cast(static_cast<__bang_bfloat16_raw>(h).x << 1U); + // saturation fixup + if (bits > (unsigned short)0xFF00U) { + // NaN + val = 0U; + } else if (f > static_cast(max_val)) { + // saturate maximum + val = max_val; + } else if (f < static_cast(min_val)) { + // saturate minimum + val = min_val; + } else { + val = static_cast(f); + } +#endif + return val; +} + +__BANG_HOSTDEVICE_BF16_DECL__ __bang_bfloat16 __ushort2bfloat16_rn__(const unsigned short int i) +{ +#if defined(__BANG_ARCH__) + __bang_bfloat16 val; + __asm__("cvt.gpr.rn.bf16.u16 \t %[dst], %[src];\n\t" + : [dst] "+r"(__BFLOAT16_TO_US(val)) + : [src] "r"(i)); + return val; +#else + const float f = static_cast(i); + return __float2bfloat16_rn__(f); +#endif +} + +__BANG_HOSTDEVICE_BF16_DECL__ bool __hisnan(const __bang_bfloat16 a) { + const __bang_bfloat16_raw hr = static_cast<__bang_bfloat16_raw>(a); + return ((hr.x & 0x7FFFU) > 0x7F80U); +} + +__BANG_HOSTDEVICE_BF16_DECL__ long long int __bfloat162ll_rz(const __bang_bfloat16 h) +{ + long long int i; +#if defined(__BANG_ARCH__) + __asm__ volatile("cvt.gpr.tz.s64.bf16 \t %[dst], %[src];\n\t" + : [dst] "=l"(i) + : [src] "r"(__BFLOAT16_TO_CUS(h))); + return i; +#else + const float f = __bfloat162float__(h); + const long long int max_val = (long long int)0x7fffffffffffffffULL; + const long long int min_val = (long long int)0x8000000000000000ULL; + const unsigned short bits = static_cast(static_cast<__bang_bfloat16_raw>(h).x << 1U); + // saturation fixup + if (bits > (unsigned short)0xFF00U) { + // NaN + i = min_val; + } else if (f >= static_cast(max_val)) { + // saturate maximum + i = max_val; + } else if (f < static_cast(min_val)) { + // saturate minimum + i = min_val; + } else { + i = static_cast(f); + } + return i; +#endif +} + +__BANG_HOSTDEVICE_BF16_DECL__ unsigned long long int __bfloat162ull_rz(const __bang_bfloat16 h) +{ + unsigned long long int i; +#if !defined(__BANG_ARCH__) || (__BANG_ARCH__ == 592) || (__BANG_ARCH__ == 506) || (__BANG_ARCH__ == 606) || (__BANG_ARCH__ == 607) + const float f = __bfloat162float__(h); + const unsigned long long int max_val = 0xffffffffffffffffULL; + const unsigned long long int min_val = 0ULL; + const unsigned short bits = static_cast(static_cast<__bang_bfloat16_raw>(h).x << 1U); + // saturation fixup + if (bits > (unsigned short)0xFF00U) { + // NaN + i = 0x8000000000000000ULL; + } else if (f >= static_cast(max_val)) { + // saturate maximum + i = max_val; + } else if (f < static_cast(min_val)) { + // saturate minimum + i = min_val; + } else { + i = static_cast(f); + } + return i; +#else + __asm__ volatile("cvt.gpr.tz.u64.bf16 \t %[dst], %[src];\n\t" + : [dst] "=l"(i) + : [src] "r"(__BFLOAT16_TO_CUS(h))); + return i; +#endif +} + +__BANG_HOSTDEVICE_BF16_DECL__ __bang_bfloat16 __ll2bfloat16_rn(const long long int i) +{ +#if defined(__BANG_ARCH__) + __bang_bfloat16 h; + __asm__ volatile("cvt.gpr.rn.bf16.s64 \t %[dst], %[src];\n\t" + : [dst] "+r"(__BFLOAT16_TO_US(h)) + : [src] "l"(i)); + return h; +#else + float f = static_cast(i); + const long long int lf = static_cast(f); + unsigned int u = __internal_float_as_uint(f); + + if ((f > 0.0f) && (lf > i)) { + u--; + } + if ((f < 0.0f) && (lf < i)) { + u--; + } + if (lf != i) { + u |= 1U; + } + + f = __internal_uint_as_float(u); + return __float2bfloat16_rn__(f); +#endif +} + +__BANG_HOSTDEVICE_BF16_DECL__ __bang_bfloat16 __ull2bfloat16_rn(const unsigned long long int i) +{ +#if defined(__BANG_ARCH__) + __bang_bfloat16 h; + __asm__ volatile("cvt.gpr.rn.bf16.u64 \t %[dst], %[src];\n\t" + : [dst] "+r"(__BFLOAT16_TO_US(h)) + : [src] "l"(i)); + return h; +#else + float f = static_cast(i); + const unsigned long long int uf = static_cast(f); + unsigned int u = __internal_float_as_uint(f); + // round up happened here + // note: no need to handle round up to f == 0x1.p64 specially + if (uf > i) { + u--; + } + if (uf != i) { + u |= 1U; + } + f = __internal_uint_as_float(u); + return __float2bfloat16_rn__(f); +#endif +} + +/****************************************************************************** +* __bang_bfloat16 comparison * +******************************************************************************/ +__BANG_HOSTDEVICE_BF16_DECL__ bool __heq(const __bang_bfloat16 a, const __bang_bfloat16 b) { + const float fa = __bfloat162float__(a); + const float fb = __bfloat162float__(b); + return (fa == fb); +} +__BANG_HOSTDEVICE_BF16_DECL__ bool __hne(const __bang_bfloat16 a, const __bang_bfloat16 b) { + const float fa = __bfloat162float__(a); + const float fb = __bfloat162float__(b); + return (fa != fb) && (!__hisnan(a) && (!__hisnan(b))); +} +__BANG_HOSTDEVICE_BF16_DECL__ bool __hle(const __bang_bfloat16 a, const __bang_bfloat16 b) { + float fa = __bfloat162float__(a); + float fb = __bfloat162float__(b); + return (fa <= fb); +} +__BANG_HOSTDEVICE_BF16_DECL__ bool __hge(const __bang_bfloat16 a, const __bang_bfloat16 b) { + float fa = __bfloat162float__(a); + float fb = __bfloat162float__(b); + return (fa >= fb); +} +__BANG_HOSTDEVICE_BF16_DECL__ bool __hlt(const __bang_bfloat16 a, const __bang_bfloat16 b) { + float fa = __bfloat162float__(a); + float fb = __bfloat162float__(b); + return (fa < fb); +} +__BANG_HOSTDEVICE_BF16_DECL__ bool __hgt(const __bang_bfloat16 a, const __bang_bfloat16 b) { + float fa = __bfloat162float__(a); + float fb = __bfloat162float__(b); + return (fa > fb); +} +__BANG_HOSTDEVICE_BF16_DECL__ bool __hequ(const __bang_bfloat16 a, const __bang_bfloat16 b) { + float fa = __bfloat162float__(a); + float fb = __bfloat162float__(b); + return (fa == fb) || (__hisnan(a)) || (__hisnan(b)); +} +__BANG_HOSTDEVICE_BF16_DECL__ bool __hneu(const __bang_bfloat16 a, const __bang_bfloat16 b) { + float fa = __bfloat162float__(a); + float fb = __bfloat162float__(b); + return (fa != fb); +} +__BANG_HOSTDEVICE_BF16_DECL__ bool __hleu(const __bang_bfloat16 a, const __bang_bfloat16 b) { + float fa = __bfloat162float__(a); + float fb = __bfloat162float__(b); + return (fa <= fb) || (__hisnan(a)) || (__hisnan(b)); +} +__BANG_HOSTDEVICE_BF16_DECL__ bool __hgeu(const __bang_bfloat16 a, const __bang_bfloat16 b) { + float fa = __bfloat162float__(a); + float fb = __bfloat162float__(b); + return (fa >= fb) || (__hisnan(a)) || (__hisnan(b)); +} +__BANG_HOSTDEVICE_BF16_DECL__ bool __hltu(const __bang_bfloat16 a, const __bang_bfloat16 b) { + float fa = __bfloat162float__(a); + float fb = __bfloat162float__(b); + return (fa < fb) || (__hisnan(a)) || (__hisnan(b)); +} +__BANG_HOSTDEVICE_BF16_DECL__ bool __hgtu(const __bang_bfloat16 a, const __bang_bfloat16 b) { + float fa = __bfloat162float__(a); + float fb = __bfloat162float__(b); + return (fa > fb) || (__hisnan(a)) || (__hisnan(b)); +} +#undef __COMPARISON_OP_BFLOAT16_MACRO +/****************************************************************************** +* __bang_bfloat16 arithmetic * +******************************************************************************/ +__BANG_HOSTDEVICE_BF16_DECL__ __bang_bfloat16 __hadd(const __bang_bfloat16 a, const __bang_bfloat16 b) { + __bang_bfloat16 val; + float fa = __bfloat162float__(a); + float fb = __bfloat162float__(b); + float fc = fa + fb; + val = __float2bfloat16__(fc); + return val; +} +__BANG_HOSTDEVICE_BF16_DECL__ __bang_bfloat16 __hsub(const __bang_bfloat16 a, const __bang_bfloat16 b) { + __bang_bfloat16 val; + float fa = __bfloat162float__(a); + float fb = __bfloat162float__(b); + float fc = fa - fb; + val = __float2bfloat16__(fc); + return val; +} +__BANG_HOSTDEVICE_BF16_DECL__ __bang_bfloat16 __hmul(const __bang_bfloat16 a, const __bang_bfloat16 b) +{ + __bang_bfloat16 val; + float fa = __bfloat162float__(a); + float fb = __bfloat162float__(b); + float fc = fa * fb; + val = __float2bfloat16__(fc); + return val; +} + +__BANG_HOSTDEVICE_BF16_DECL__ __bang_bfloat16 __hdiv(const __bang_bfloat16 a, const __bang_bfloat16 b) { + __bang_bfloat16 val; + float fa = __bfloat162float__(a); + float fb = __bfloat162float__(b); + float fc = fa / fb; + val = __float2bfloat16__(fc); + return val; +} + +__BANG_HOSTDEVICE_BF16_DECL__ __bang_bfloat16 __hneg(const __bang_bfloat16 a) { + float fa = __bfloat162float__(a); + return __float2bfloat16__(-fa); +} +#endif /* defined(__cplusplus) */ + +#undef __BANG_HOSTDEVICE_BF16_DECL__ + +#if defined(__cplusplus) && !defined(BANG_NO_BFLOAT16) +typedef __bang_bfloat16 bfloat16_t; +typedef __bang_bfloat16_raw bfloat16_raw_t; +#endif /* defined(__cplusplus) && !defined(BANG_NO_BFLOAT16) */ + +#if defined(__CPP_VERSION_AT_LEAST_11_BF16) +#undef __CPP_VERSION_AT_LEAST_11_BF16 +#endif /* defined(__CPP_VERSION_AT_LEAST_11_BF16) */ + +#endif /* end of include guard: __BANG_BF16_H__ */ diff --git a/backends/mlu/third_party/bang_host/bang_fp16.h b/backends/mlu/third_party/bang_host/bang_fp16.h new file mode 100644 index 0000000..8ffc9bc --- /dev/null +++ b/backends/mlu/third_party/bang_host/bang_fp16.h @@ -0,0 +1,1150 @@ +/*===-----------------------------------------------------------------------=== + * + * Copyright (C) [2022] by Cambricon, Inc. + * + * Permission is hereby granted, free of charge, to any person obtaining a copy + * of this software and associated documentation files (the "Software"), to deal + * in the Software without restriction, including without limitation the rights + * to use, copy, modify, merge, publish, distribute, sublicense, and/or sell + * copies of the Software, and to permit persons to whom the Software is + * furnished to do so, subject to the following conditions: + * + * The above copyright notice and this permission notice shall be included in + * all copies or substantial portions of the Software. + * + * THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR + * IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, + * FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE + * AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER + * LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, + * OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN + * THE SOFTWARE. + * + *===-----------------------------------------------------------------------=== + */ + +#ifndef __BANG_FP16_H__ +#define __BANG_FP16_H__ + +#if defined(__BANGCC__) +#define __BANG_HOSTDEVICE_FP16_DECL__ static __mlu_host__ __mlu_func__ +#define __BANG_HOSTDEVICE__ __mlu_host__ __mlu_func__ inline +#else /* !defined(__BANGCC__) */ +#if defined(__GNUC__) +#define __BANG_HOSTDEVICE_FP16_DECL__ static __attribute__((unused)) +#else +#define __BANG_HOSTDEVICE_FP16_DECL__ static +#endif /* defined(__GNUC__) */ +#define __BANG_HOSTDEVICE__ inline +#endif /* defined(__BANGCC_) */ + +#if !defined(_MSC_VER) && __cplusplus >= 201103L +# define __CPP_VERSION_AT_LEAST_11_FP16 +#elif _MSC_FULL_VER >= 190024210 && _MSVC_LANG >= 201103L +# define __CPP_VERSION_AT_LEAST_11_FP16 +#endif + +/* Set up structure-alignment attribute */ +/* Define alignment macro based on compiler type (cannot assume C11 "_Alignas" is available) */ +#if __cplusplus >= 201103L +#define __BANG_ALIGN__(n) alignas(n) /* C++11 kindly gives us a keyword for this */ +#else /* !defined(__CPP_VERSION_AT_LEAST_11_FP16)*/ +#if defined(__GNUC__) +#define __BANG_ALIGN__(n) __attribute__ ((aligned(n))) +#elif defined(_MSC_VER) +#define __BANG_ALIGN__(n) __declspec(align(n)) +#else +#define __BANG_ALIGN__(n) +#endif /* defined(__GNUC__) */ +#endif /* defined(__CPP_VERSION_AT_LEAST_11_FP16) */ + +// define __BANG_FP16_CONSTEXPR__ in order to +// use constexpr where possible, with supporting C++ dialects +// undef after use +#if (defined __CPP_VERSION_AT_LEAST_11_FP16) +#define __BANG_FP16_CONSTEXPR__ constexpr +#else +#define __BANG_FP16_CONSTEXPR__ +#endif + +typedef struct __BANG_ALIGN__(2) { + unsigned short x; +} __half_raw; + +/* Hide GCC member initialization list warnings because of host/device in-function init requirement */ +#if defined(__GNUC__) +#if __GNUC__ > 4 || (__GNUC__ == 4 && __GNUC_MINOR__ >= 6) +#pragma GCC diagnostic push +#pragma GCC diagnostic ignored "-Wstrict-aliasing" +#pragma GCC diagnostic ignored "-Weffc++" +#endif /* __GNUC__ > 4 || (__GNUC__ == 4 && __GNUC_MINOR__ >= 6) */ +#endif /* defined(__GNUC__) */ + +/* class' : multiple assignment operators specified + The class has multiple assignment operators of a single type. This warning is informational */ +#if defined(_MSC_VER) && _MSC_VER >= 1500 +#pragma warning( push ) +#pragma warning( disable:4522 ) +#endif /* defined(_MSC_VER) && _MSC_VER >= 1500 */ + +// forward-declaration of bfloat type to be used in converting constructor +struct __bang_bfloat16; + +/* Macros to allow half to be used by inline assembly */ +#define __HALF_TO_US(var) *(reinterpret_cast(&(var))) +#define __HALF_TO_CUS(var) *(reinterpret_cast(&(var))) + +struct __BANG_ALIGN__(2) __half { +protected: + /** + * Protected storage variable contains the bits of floating-point data. + */ + unsigned short __x; + +public: + /** + * Constructor by default. + */ +#if defined(__CPP_VERSION_AT_LEAST_11_FP16) + __half() = default; +#else + __BANG_HOSTDEVICE__ __half(); +#endif /* defined(__CPP_VERSION_AT_LEAST_11_FP16) */ + + /* Convert to/from __half_raw */ + /** + * Constructor from \p __half_raw. + */ + __BANG_HOSTDEVICE__ __BANG_FP16_CONSTEXPR__ __half(const __half_raw &hr) : __x(hr.x) { } + /** + * Assignment operator from \p __half_raw. + */ + __BANG_HOSTDEVICE__ __half &operator=(const __half_raw &hr); + /** + * Assignment operator from \p __half_raw to \p volatile \p __half. + */ + __BANG_HOSTDEVICE__ volatile __half &operator=(const __half_raw &hr) volatile; + /** + * Assignment operator from \p volatile \p __half_raw to \p volatile \p __half. + */ + __BANG_HOSTDEVICE__ volatile __half &operator=(const volatile __half_raw &hr) volatile; + /** + * Type cast to \p __half_raw operator. + */ + __BANG_HOSTDEVICE__ operator __half_raw() const; + /** + * Type cast to \p __half_raw operator with \p volatile input. + */ + __BANG_HOSTDEVICE__ operator __half_raw() const volatile; + /** + * Construct \p __half from \p __nv_bfloat16 input using default round-to-nearest-even rounding mode. + * Need to include the header file \p bang_bf16.h + */ + __BANG_HOSTDEVICE__ __half(const __bang_bfloat16 f); //forward declaration only, implemented in bang_bf16.hpp + /* Construct from float/double */ + /** + * Construct \p __half from \p float input using default round-to-nearest-even rounding mode. + */ + __BANG_HOSTDEVICE__ __half(const float f); + /** + * Construct \p __half from \p double input using default round-to-nearest-even rounding mode. + */ + __BANG_HOSTDEVICE__ __half(const double f); + /** + * Type cast to \p float operator. + */ + __BANG_HOSTDEVICE__ operator float() const; + /** + * Type cast to \p __half assignment operator from \p float input using default round-to-nearest-even rounding mode. + */ + __BANG_HOSTDEVICE__ __half &operator=(const float f); + + /* We omit "cast to double" operator, so as to not be ambiguous about up-cast */ + /** + * Type cast to \p __half assignment operator from \p double input using default round-to-nearest-even rounding mode. + */ + __BANG_HOSTDEVICE__ __half &operator=(const double f); + + /* Allow automatic construction from types supported natively in hardware */ + /* Note we do avoid constructor init-list because of special host/device compilation rules */ + + /** + * Construct \p __half from \p short integer input using default round-to-nearest-even rounding mode. + */ + __BANG_HOSTDEVICE__ __half(const short val); + /** + * Construct \p __half from \p unsigned \p short integer input using default round-to-nearest-even rounding mode. + */ + __BANG_HOSTDEVICE__ __half(const unsigned short val); + /** + * Construct \p __half from \p int input using default round-to-nearest-even rounding mode. + */ + __BANG_HOSTDEVICE__ __half(const int val); + /** + * Construct \p __half from \p unsigned \p int input using default round-to-nearest-even rounding mode. + */ + __BANG_HOSTDEVICE__ __half(const unsigned int val); + /** + * Construct \p __half from \p long input using default round-to-nearest-even rounding mode. + */ + __BANG_HOSTDEVICE__ __half(const long val); + /** + * Construct \p __half from \p unsigned \p long input using default round-to-nearest-even rounding mode. + */ + __BANG_HOSTDEVICE__ __half(const unsigned long val); + + /** + * Construct \p __half from \p long \p long input using default round-to-nearest-even rounding mode. + */ + __BANG_HOSTDEVICE__ __half(const long long val); + /** + * Construct \p __half from \p unsigned \p long \p long input using default round-to-nearest-even rounding mode. + */ + __BANG_HOSTDEVICE__ __half(const unsigned long long val); + + /* Allow automatic casts to supported builtin types, matching all that are permitted with float */ + + /** + * Conversion operator to \p signed \p char data type. + * Using round-toward-zero rounding mode. + * + * See __half2char_rz(__half) for further details + */ + __BANG_HOSTDEVICE__ operator signed char() const; + /** + * Conversion operator to \p unsigned \p char data type. + * Using round-toward-zero rounding mode. + * + * See __half2uchar_rz(__half) for further details + */ + __BANG_HOSTDEVICE__ operator unsigned char() const; + /** + * Conversion operator to an implementation defined \p char data type. + * Using round-toward-zero rounding mode. + * + * Detects signedness of the \p char type and proceeds accordingly, see + * further details in signed and unsigned char operators. + */ + __BANG_HOSTDEVICE__ operator char() const; + /** + * Conversion operator to \p short data type. + * Using round-toward-zero rounding mode. + * + * See __half2short_rz(__half) for further details + */ + __BANG_HOSTDEVICE__ operator short() const; + /** + * Conversion operator to \p unsigned \p short data type. + * Using round-toward-zero rounding mode. + * + * See __half2ushort_rz(__half) for further details + */ + __BANG_HOSTDEVICE__ operator unsigned short() const; + /** + * Conversion operator to \p int data type. + * Using round-toward-zero rounding mode. + * + * See __half2int_rz(__half) for further details + */ + __BANG_HOSTDEVICE__ operator int() const; + /** + * Conversion operator to \p unsigned \p int data type. + * Using round-toward-zero rounding mode. + * + * See __half2uint_rz(__half) for further details + */ + __BANG_HOSTDEVICE__ operator unsigned int() const; + /** + * Conversion operator to \p long data type. + * Using round-toward-zero rounding mode. + */ + __BANG_HOSTDEVICE__ operator long() const; + /** + * Conversion operator to \p unsigned \p long data type. + * Using round-toward-zero rounding mode. + */ + __BANG_HOSTDEVICE__ operator unsigned long() const; + /** + * Conversion operator to \p long \p long data type. + * Using round-toward-zero rounding mode. + * + * See __half2ll_rz(__half) for further details + */ + __BANG_HOSTDEVICE__ operator long long() const; + /** + * Conversion operator to \p unsigned \p long \p long data type. + * Using round-toward-zero rounding mode. + * + * See __half2ull_rz(__half) for further details + */ + __BANG_HOSTDEVICE__ operator unsigned long long() const; + /** + * Type cast from \p short assignment operator, using default round-to-nearest-even rounding mode. + */ + __BANG_HOSTDEVICE__ __half &operator=(const short val); + /** + * Type cast from \p unsigned \p short assignment operator, using default round-to-nearest-even rounding mode. + */ + __BANG_HOSTDEVICE__ __half &operator=(const unsigned short val); + /** + * Type cast from \p int assignment operator, using default round-to-nearest-even rounding mode. + */ + __BANG_HOSTDEVICE__ __half &operator=(const int val); + /** + * Type cast from \p unsigned \p int assignment operator, using default round-to-nearest-even rounding mode. + */ + __BANG_HOSTDEVICE__ __half &operator=(const unsigned int val); + /** + * Type cast from \p long \p long assignment operator, using default round-to-nearest-even rounding mode. + */ + __BANG_HOSTDEVICE__ __half &operator=(const long long val); + /** + * Type cast from \p long assignment operator, using default round-to-nearest-even rounding mode. + */ + __BANG_HOSTDEVICE__ __half &operator=(const long val); + /** + * Type cast from \p unsigned \p long assignment operator, using default round-to-nearest-even rounding mode. + */ + __BANG_HOSTDEVICE__ __half &operator=(const unsigned long val); + /** + * Type cast from \p unsigned \p long \p long assignment operator, using default round-to-nearest-even rounding mode. + */ + __BANG_HOSTDEVICE__ __half &operator=(const unsigned long long val); + /** + * Type cast from \p __bang_bfloat16 assignment operator, using default round-to-nearest-even rounding mode. + */ + __BANG_HOSTDEVICE__ __half &operator=(const __bang_bfloat16 f); + /** + * Conversion operator to \p bool data type. + * +0 and -0 inputs convert to \p false. + * Non-zero inputs convert to \p true. + */ + __BANG_HOSTDEVICE__ __BANG_FP16_CONSTEXPR__ operator bool() const { return (__x & 0x7FFFU) != 0U; } +}; + +__BANG_HOSTDEVICE_FP16_DECL__ __half __float2half__(const float a); +__BANG_HOSTDEVICE_FP16_DECL__ float __half2float(const __half a); +__BANG_HOSTDEVICE_FP16_DECL__ short int __half2short_rz(const __half h); + +__BANG_HOSTDEVICE_FP16_DECL__ long long int __half2ll_rz(const __half h); +__BANG_HOSTDEVICE_FP16_DECL__ unsigned long long int __half2ull_rz(const __half h); +__BANG_HOSTDEVICE_FP16_DECL__ __half __short2half_rn__(const short int i); +__BANG_HOSTDEVICE_FP16_DECL__ __half __ushort2half_rn__(const unsigned short int i); + +__BANG_HOSTDEVICE_FP16_DECL__ __half __int2half_rn__(const int i); +__BANG_HOSTDEVICE_FP16_DECL__ __half __uint2half_rn__(const unsigned int i); +__BANG_HOSTDEVICE_FP16_DECL__ __half __ll2half_rn(const long long int i); +__BANG_HOSTDEVICE_FP16_DECL__ __half __ull2half_rn(const unsigned long long int i); +__BANG_HOSTDEVICE_FP16_DECL__ unsigned int __half2uint_rz(const __half h); +__BANG_HOSTDEVICE_FP16_DECL__ int __half2int_rz(const __half h); +__BANG_HOSTDEVICE_FP16_DECL__ unsigned short int __half2ushort_rz(const __half h); +__BANG_HOSTDEVICE_FP16_DECL__ unsigned char __half2uchar_rz(const __half h); +__BANG_HOSTDEVICE_FP16_DECL__ signed char __half2char_rz(const __half h); + +__BANG_HOSTDEVICE_FP16_DECL__ __half __hneg(const __half a); +__BANG_HOSTDEVICE_FP16_DECL__ bool __heq(const __half a, const __half b); +__BANG_HOSTDEVICE_FP16_DECL__ bool __hne(const __half a, const __half b); +__BANG_HOSTDEVICE_FP16_DECL__ bool __hle(const __half a, const __half b); +__BANG_HOSTDEVICE_FP16_DECL__ bool __hge(const __half a, const __half b); +__BANG_HOSTDEVICE_FP16_DECL__ bool __hlt(const __half a, const __half b); +__BANG_HOSTDEVICE_FP16_DECL__ bool __hgt(const __half a, const __half b); +__BANG_HOSTDEVICE_FP16_DECL__ bool __hneu(const __half a, const __half b); + +__BANG_HOSTDEVICE_FP16_DECL__ __half __hadd(const __half a, const __half b); +__BANG_HOSTDEVICE_FP16_DECL__ __half __hsub(const __half a, const __half b); +__BANG_HOSTDEVICE_FP16_DECL__ __half __hmul(const __half a, const __half b); +__BANG_HOSTDEVICE_FP16_DECL__ __half __hdiv(const __half a, const __half b); + +__BANG_HOSTDEVICE_FP16_DECL__ bool __hisnan(const __half a); + +#if defined(__CPP_VERSION_AT_LEAST_11_FP16) +#else + __BANG_HOSTDEVICE__ __half::__half() { } +#endif /* defined(__CPP_VERSION_AT_LEAST_11_FP16) */ +__BANG_HOSTDEVICE__ __half &__half::operator=(const __half_raw &hr) { + __x = hr.x; + return *this; +} +__BANG_HOSTDEVICE__ volatile __half & __half::operator=(const __half_raw &hr) volatile { __x = hr.x; return *this; } +__BANG_HOSTDEVICE__ volatile __half & __half::operator=(const volatile __half_raw &hr) volatile { __x = hr.x; return *this; } +__BANG_HOSTDEVICE__ __half::operator __half_raw() const { __half_raw ret; ret.x = __x; return ret; } +__BANG_HOSTDEVICE__ __half::operator __half_raw() const volatile { __half_raw ret; ret.x = __x; return ret; } + +__BANG_HOSTDEVICE__ __half::__half(const float f) { __x = __float2half__(f).__x; } +__BANG_HOSTDEVICE__ __half::operator float() const { return __half2float(*this); } +__BANG_HOSTDEVICE__ __half &__half::operator=(const float f) { __x = __float2half__(f).__x; return *this; } + +__BANG_HOSTDEVICE__ __half::__half(const short val) { __x = __short2half_rn__(val).__x; } +__BANG_HOSTDEVICE__ __half::__half(const unsigned short val) { __x = __ushort2half_rn__(val).__x; } +__BANG_HOSTDEVICE__ __half::__half(const int val) { __x = __int2half_rn__(val).__x; } +__BANG_HOSTDEVICE__ __half::__half(const unsigned int val) { __x = __uint2half_rn__(val).__x; } + +__BANG_HOSTDEVICE__ __half::__half(const long val) { + +/* Suppress VS warning: warning C4127: conditional expression is constant */ +#if defined(_MSC_VER) +#pragma warning(disable : 4127) +#endif /* _MSC_VER */ + if (sizeof(long) == sizeof(long long)) +#if defined(_MSC_VER) +#pragma warning(default : 4127) +#endif /* _MSC_VER */ + + { + __x = __ll2half_rn(static_cast(val)).__x; + } else { + __x = __int2half_rn__(static_cast(val)).__x; + } +} + +__BANG_HOSTDEVICE__ __half::__half(const unsigned long val) { + +/* Suppress VS warning: warning C4127: conditional expression is + * constant */ +#if defined(_MSC_VER) +#pragma warning(disable : 4127) +#endif /* _MSC_VER */ +if (sizeof(unsigned long) == sizeof(unsigned long long)) +#if defined(_MSC_VER) +#pragma warning(default : 4127) +#endif /* _MSC_VER */ + +{ + __x = __ull2half_rn(static_cast(val)).__x; +} else { + __x = __uint2half_rn__(static_cast(val)).__x; +} +} + +__BANG_HOSTDEVICE__ __half::__half(const long long val) { __x = __ll2half_rn(val).__x; } +__BANG_HOSTDEVICE__ __half::__half(const unsigned long long val) { __x = __ull2half_rn(val).__x; } +__BANG_HOSTDEVICE__ __half::operator signed char() const { return __half2char_rz(*this); } +__BANG_HOSTDEVICE__ __half::operator unsigned char() const { return __half2uchar_rz(*this); } +__BANG_HOSTDEVICE__ __half::operator char() const { + char value; + /* Suppress VS warning: warning C4127: conditional expression is constant + */ +#if defined(_MSC_VER) +#pragma warning(disable : 4127) +#endif /* _MSC_VER */ + if (((char)-1) < (char)0) +#if defined(_MSC_VER) +#pragma warning(default : 4127) +#endif /* _MSC_VER */ + { + value = static_cast(__half2char_rz(*this)); + } else { + value = static_cast(__half2uchar_rz(*this)); + } + return value; +} + +__BANG_HOSTDEVICE__ __half::operator short() const { return __half2short_rz(*this); } +__BANG_HOSTDEVICE__ __half::operator unsigned short() const { return __half2ushort_rz(*this); } +__BANG_HOSTDEVICE__ __half::operator int() const { return __half2int_rz(*this); } +__BANG_HOSTDEVICE__ __half::operator unsigned int() const { return __half2uint_rz(*this); } + +__BANG_HOSTDEVICE__ __half::operator long() const { + long retval; + /* Suppress VS warning: warning C4127: conditional expression is constant + */ +#if defined(_MSC_VER) +#pragma warning(disable : 4127) +#endif /* _MSC_VER */ + if (sizeof(long) == sizeof(long long)) +#if defined(_MSC_VER) +#pragma warning(default : 4127) +#endif /* _MSC_VER */ + { + retval = static_cast(__half2ll_rz(*this)); + } else { + retval = static_cast(__half2int_rz(*this)); + } + return retval; +} + +__BANG_HOSTDEVICE__ __half::operator unsigned long() const { + unsigned long retval; +/* Suppress VS warning: warning C4127: conditional expression is constant */ +#if defined(_MSC_VER) +#pragma warning(disable : 4127) +#endif /* _MSC_VER */ + if (sizeof(unsigned long) == sizeof(unsigned long long)) +#if defined(_MSC_VER) +#pragma warning(default : 4127) +#endif /* _MSC_VER */ + { + retval = static_cast(__half2ull_rz(*this)); + } else { + retval = static_cast(__half2uint_rz(*this)); + } + return retval; +} + +__BANG_HOSTDEVICE__ __half::operator long long() const { return __half2ll_rz(*this); } +__BANG_HOSTDEVICE__ __half::operator unsigned long long() const { return __half2ull_rz(*this); } +__BANG_HOSTDEVICE__ __half &__half::operator=(const short val) { __x = __short2half_rn__(val).__x; return *this; } +__BANG_HOSTDEVICE__ __half &__half::operator=(const unsigned short val) { __x = __ushort2half_rn__(val).__x; return *this; } +__BANG_HOSTDEVICE__ __half &__half::operator=(const int val) { __x = __int2half_rn__(val).__x; return *this; } +__BANG_HOSTDEVICE__ __half &__half::operator=(const unsigned int val) { __x = __uint2half_rn__(val).__x; return *this; } +__BANG_HOSTDEVICE__ __half &__half::operator=(const long long val) { __x = __ll2half_rn(val).__x; return *this; } +__BANG_HOSTDEVICE__ __half &__half::operator=(const long val) { __x = __half(val).__x; return *this; } +__BANG_HOSTDEVICE__ __half &__half::operator=(const unsigned long val) { __x = __half(val).__x; return *this; } +__BANG_HOSTDEVICE__ __half &__half::operator=(const unsigned long long val) { __x = __ull2half_rn(val).__x; return *this; } + +__BANG_HOSTDEVICE__ __half operator+(const __half &lh, const __half &rh) { return __hadd(lh, rh); } +__BANG_HOSTDEVICE__ __half operator-(const __half &lh, const __half &rh) { return __hsub(lh, rh); } +__BANG_HOSTDEVICE__ __half operator*(const __half &lh, const __half &rh) { return __hmul(lh, rh); } +__BANG_HOSTDEVICE__ __half operator/(const __half &lh, const __half &rh) { return __hdiv(lh, rh); } +__BANG_HOSTDEVICE__ __half &operator+=(__half &lh, const __half &rh) { lh = __hadd(lh, rh); return lh; } +__BANG_HOSTDEVICE__ __half &operator-=(__half &lh, const __half &rh) { lh = __hsub(lh, rh); return lh; } +__BANG_HOSTDEVICE__ __half &operator*=(__half &lh, const __half &rh) { lh = __hmul(lh, rh); return lh; } +__BANG_HOSTDEVICE__ __half &operator/=(__half &lh, const __half &rh) { lh = __hdiv(lh, rh); return lh; } +__BANG_HOSTDEVICE__ __half &operator++(__half &h) { __half_raw one; one.x = 0x3C00U; h += one; return h; } +__BANG_HOSTDEVICE__ __half &operator--(__half &h) { __half_raw one; one.x = 0x3C00U; h -= one; return h; } +__BANG_HOSTDEVICE__ __half operator++(__half &h, const int ignored) { + // ignored on purpose. Parameter only needed to distinguish the function declaration from other types of operators. + static_cast(ignored); + + const __half ret = h; + __half_raw one; + one.x = 0x3C00U; + h += one; + return ret; +} + +__BANG_HOSTDEVICE__ __half operator--(__half &h, const int ignored) +{ + // ignored on purpose. Parameter only needed to distinguish the function declaration from other types of operators. + static_cast(ignored); + + const __half ret = h; + __half_raw one; + one.x = 0x3C00U; + h -= one; + return ret; +} + +__BANG_HOSTDEVICE__ __half operator+(const __half &h) { return h; } +__BANG_HOSTDEVICE__ __half operator-(const __half &h) { return __hneg(h); } +__BANG_HOSTDEVICE__ bool operator==(const __half &lh, const __half &rh) { return __heq(lh, rh); } +__BANG_HOSTDEVICE__ bool operator!=(const __half &lh, const __half &rh) { return __hneu(lh, rh); } +__BANG_HOSTDEVICE__ bool operator> (const __half &lh, const __half &rh) { return __hgt(lh, rh); } +__BANG_HOSTDEVICE__ bool operator< (const __half &lh, const __half &rh) { return __hlt(lh, rh); } +__BANG_HOSTDEVICE__ bool operator>=(const __half &lh, const __half &rh) { return __hge(lh, rh); } +__BANG_HOSTDEVICE__ bool operator<=(const __half &lh, const __half &rh) { return __hle(lh, rh); } + +/* Restore warning for multiple assignment operators */ +#if defined(_MSC_VER) && _MSC_VER >= 1500 +#pragma warning( pop ) +#endif /* defined(_MSC_VER) && _MSC_VER >= 1500 */ + +/* Restore -Weffc++ warnings from here on */ +#if defined(__GNUC__) +#if __GNUC__ > 4 || (__GNUC__ == 4 && __GNUC_MINOR__ >= 6) +#pragma GCC diagnostic pop +#endif /* __GNUC__ > 4 || (__GNUC__ == 4 && __GNUC_MINOR__ >= 6) */ +#endif /* defined(__GNUC__) */ + +#undef __BANG_HOSTDEVICE__ +#undef __BANG_ALIGN__ + +static inline unsigned short __internal_float2half(const float f, unsigned int &sign, unsigned int &remainder) +{ + unsigned int u; + unsigned int result; + unsigned int x = *(reinterpret_cast(&(f))); + u = (x & 0x7fffffffU); + sign = ((x >> 16U) & 0x8000U); + // NaN/+Inf/-Inf + if (u >= 0x7f800000U) { + remainder = 0U; + result = ((u == 0x7f800000U) ? (sign | 0x7c00U) : 0x7fffU); + } else if (u > 0x477fefffU) { // Overflows + remainder = 0x80000000U; + result = (sign | 0x7bffU); + } else if (u >= 0x38800000U) { // Normal numbers + remainder = u << 19U; + u -= 0x38000000U; + result = (sign | (u >> 13U)); + } else if (u < 0x33000001U) { // +0/-0 + remainder = u; + result = sign; + } else { // Denormal numbers + const unsigned int exponent = u >> 23U; + const unsigned int shift = 0x7eU - exponent; + unsigned int mantissa = (u & 0x7fffffU); + mantissa |= 0x800000U; + remainder = mantissa << (32U - shift); + result = (sign | (mantissa >> shift)); + result &= 0x0000FFFFU; + } + return static_cast(result); +} + +static inline float __internal_half2float(const unsigned short h) { + unsigned int sign = ((static_cast(h) >> 15U) & 1U); + unsigned int exponent = ((static_cast(h) >> 10U) & 0x1fU); + unsigned int mantissa = ((static_cast(h) & 0x3ffU) << 13U); + float f; + if (exponent == 0x1fU) { /* NaN or Inf */ + /* discard sign of a NaN */ + sign = ((mantissa != 0U) ? (sign >> 1U) : sign); + mantissa = ((mantissa != 0U) ? 0x7fffffU : 0U); + exponent = 0xffU; + } else if (exponent == 0U) { /* Denorm or Zero */ + if (mantissa != 0U) { + unsigned int msb; + exponent = 0x71U; + do { + msb = (mantissa & 0x400000U); + mantissa <<= 1U; /* normalize */ + --exponent; + } while (msb == 0U); + mantissa &= 0x7fffffU; /* 1.mantissa is implicit */ + } + } else { + exponent += 0x70U; + } + const unsigned int u = ((sign << 31U) | (exponent << 23U) | mantissa); + f = *(reinterpret_cast(&(u))); + return f; +} + +__BANG_HOSTDEVICE_FP16_DECL__ __half __float2half__(const float a) { + __half val; +#ifdef __BANG_ARCH__ + asm("cvt.gpr.rn.f16.f32 %0, %1;" : "+r"(__HALF_TO_US(val)) : "r"(a)); +#else + __half_raw r; + unsigned int sign = 0U; + unsigned int remainder = 0U; + r.x = __internal_float2half(a, sign, remainder); + if ((remainder > 0x80000000U) || + ((remainder == 0x80000000U) && ((r.x & 0x1U) != 0U))) { + r.x++; + } + val = r; +#endif + return val; +} + +__BANG_HOSTDEVICE_FP16_DECL__ float __half2float(const __half a) { + float val; +#ifdef __BANG_ARCH__ + asm("cvt.gpr.f32.f16 %0, %1;" : "+r"(val) : "r"(__HALF_TO_CUS(a))); +#else + val = __internal_half2float(static_cast<__half_raw>(a).x); +#endif + return val; +} + +__BANG_HOSTDEVICE_FP16_DECL__ __half __float2half_rn__(const float a) { + __half val; +#ifdef __BANG_ARCH__ + asm("cvt.gpr.rn.f16.f32 %0, %1;" : "+r"(__HALF_TO_US(val)) : "r"(a)); +#else + __half_raw r; + unsigned int sign = 0U; + unsigned int remainder = 0U; + r.x = __internal_float2half(a, sign, remainder); + if ((remainder > 0x80000000U) || + ((remainder == 0x80000000U) && ((r.x & 0x1U) != 0U))) { + r.x++; + } + val = r; +#endif + return val; +} + +__BANG_HOSTDEVICE_FP16_DECL__ __half __short2half_rn__(const short int i) { + __half h; +#ifdef __BANG_ARCH__ + asm("cvt.gpr.rn.f16.s16 %0, %1;" : "+r"(__HALF_TO_US(h)) : "r"(i)); +#else + const float f = static_cast(i); + h = __float2half_rn__(f); +#endif + return h; +} + +__BANG_HOSTDEVICE_FP16_DECL__ __half __int2half_rn__(const int i) { + __half h; +#ifdef __BANG_ARCH__ + asm("cvt.gpr.rn.f16.s32 %0, %1;" : "+r"(__HALF_TO_US(h)) : "r"(i)); +#else + // double-rounding is not a problem here: if integer + // has more than 24 bits, it is already too large to + // be represented in half precision, and result will + // be infinity. + const float f = static_cast(i); + h = __float2half_rn__(f); +#endif + return h; +} + +__BANG_HOSTDEVICE_FP16_DECL__ __half __uint2half_rn__(const unsigned int i) { + __half h; +#ifdef __BANG_ARCH__ + asm("cvt.gpr.rn.f16.u32 %0, %1;" : "+r"(__HALF_TO_US(h)) : "r"(i)); +#else + // double-rounding is not a problem here: if integer + // has more than 24 bits, it is already too large to + // be represented in half precision, and result will + // be infinity. + const float f = static_cast(i); + h = __float2half_rn__(f); +#endif + return h; +} + +__BANG_HOSTDEVICE_FP16_DECL__ long long int __half2ll_rz(const __half h) { + long long int i; +#ifdef __BANG_ARCH__ + asm("cvt.gpr.tz.s64.f16 %0, %1;" : "+r"(i) : "r"(__HALF_TO_CUS(h))); +#else + const float f = __half2float(h); + const long long int max_val = (long long int)0x7fffffffffffffffULL; + const long long int min_val = (long long int)0x8000000000000000ULL; + const unsigned short bits = + static_cast(static_cast<__half_raw>(h).x << 1U); + // saturation fixup + if (bits > (unsigned short)0xF800U) { + // NaN + i = min_val; + } else if (f > static_cast(max_val)) { + // saturate maximum + i = max_val; + } else if (f < static_cast(min_val)) { + // saturate minimum + i = min_val; + } else { + // normal value, conversion is well-defined + i = static_cast(f); + } +#endif + return i; +} + +__BANG_HOSTDEVICE_FP16_DECL__ unsigned long long int __half2ull_rz(const __half h) { + unsigned long long int i; + +#ifdef __BANG_ARCH__ + asm("cvt.gpr.tz.u64.f16 %0, %1;" : "+r"(i) : "r"(__HALF_TO_CUS(h))); +#else + const float f = __half2float(h); + const unsigned long long int max_val = 0xffffffffffffffffULL; + const unsigned long long int min_val = 0ULL; + const unsigned short bits = + static_cast(static_cast<__half_raw>(h).x << 1U); + // saturation fixup + if (bits > (unsigned short)0xF800U) { + // NaN + i = 0x8000000000000000ULL; + } else if (f > static_cast(max_val)) { + // saturate maximum + i = max_val; + } else if (f < static_cast(min_val)) { + // saturate minimum + i = min_val; + } else { + // normal value, conversion is well-defined + i = static_cast(f); + } +#endif + return i; +} + +__BANG_HOSTDEVICE_FP16_DECL__ __half __ushort2half_rn__(const unsigned short int i) { + __half h; +#ifdef __BANG_ARCH__ + asm("cvt.gpr.rn.f16.u16 %0, %1;" : "+r"(__HALF_TO_US(h)) : "r"(i)); +#else + const float f = static_cast(i); + h = __float2half_rn__(f); +#endif + + return h; +} + +__BANG_HOSTDEVICE_FP16_DECL__ __half __ll2half_rn(const long long int i) { + __half h; +#ifdef __BANG_ARCH__ + asm("cvt.gpr.rn.f16.s64 %0, %1;" : "+r"(__HALF_TO_US(h)) : "r"(i)); +#else + // double-rounding is not a problem here: if integer + // has more than 24 bits, it is already too large to + // be represented in half precision, and result will + // be infinity. + const float f = static_cast(i); + h = __float2half_rn__(f); +#endif + return h; +} + +__BANG_HOSTDEVICE_FP16_DECL__ __half __ull2half_rn(const unsigned long long int i) { + __half h; +#ifdef __BANG_ARCH__ + asm("cvt.gpr.rn.f16.u64 %0, %1;" : "+r"(__HALF_TO_US(h)) : "r"(i)); +#else + // double-rounding is not a problem here: if integer + // has more than 24 bits, it is already too large to + // be represented in half precision, and result will + // be infinity. + const float f = static_cast(i); + h = __float2half_rn__(f); +#endif + return h; +} + +__BANG_HOSTDEVICE_FP16_DECL__ unsigned int __half2uint_rz(const __half h) { + unsigned int i; + +#ifdef __BANG_ARCH__ + asm("cvt.gpr.tz.u32.f16 %0, %1;" : "=r"(i) : "r"(__HALF_TO_CUS(h))); +#else + const float f = __half2float(h); + const unsigned int max_val = 0xffffffffU; + const unsigned int min_val = 0U; + const unsigned short bits = + static_cast(static_cast<__half_raw>(h).x << 1U); + // saturation fixup + if (bits > (unsigned short)0xF800U) { + // NaN + i = 0U; + } else if (f > static_cast(max_val)) { + // saturate maximum + i = max_val; + } else if (f < static_cast(min_val)) { + // saturate minimum + i = min_val; + } else { + // normal value, conversion is well-defined + i = static_cast(f); + } +#endif + return i; +} + +__BANG_HOSTDEVICE_FP16_DECL__ int __half2int_rz(const __half h) { + int i; +#ifdef __BANG_ARCH__ + asm("cvt.gpr.tz.s32.f16 %0, %1;" : "=r"(i) : "r"(__HALF_TO_CUS(h))); +#else + const float f = __half2float(h); + const int max_val = (int)0x7fffffffU; + const int min_val = (int)0x80000000U; + const unsigned short bits = + static_cast(static_cast<__half_raw>(h).x << 1U); + // saturation fixup + if (bits > (unsigned short)0xF800U) { + // NaN + i = 0; + } else if (f > static_cast(max_val)) { + // saturate maximum + i = max_val; + } else if (f < static_cast(min_val)) { + // saturate minimum + i = min_val; + } else { + // normal value, conversion is well-defined + i = static_cast(f); + } +#endif + return i; +} + +__BANG_HOSTDEVICE_FP16_DECL__ unsigned short int __half2ushort_rz(const __half h) { + unsigned short int i; + +#ifdef __BANG_ARCH__ + asm("cvt.gpr.tz.u16.f16 %0, %1;" : "+r"(i) : "r"(__HALF_TO_CUS(h))); +#else + const float f = __half2float(h); + const unsigned short int max_val = 0xffffU; + const unsigned short int min_val = 0U; + const unsigned short bits = + static_cast(static_cast<__half_raw>(h).x << 1U); + // saturation fixup + if (bits > (unsigned short)0xF800U) { + // NaN + i = 0U; + } else if (f > static_cast(max_val)) { + // saturate maximum + i = max_val; + } else if (f < static_cast(min_val)) { + // saturate minimum + i = min_val; + } else { + // normal value, conversion is well-defined + i = static_cast(f); + } +#endif + return i; +} + +__BANG_HOSTDEVICE_FP16_DECL__ unsigned char __half2uchar_rz(const __half h) { + unsigned char i; + +#ifdef __BANG_ARCH__ + unsigned int tmp; + asm("cvt.gpr.tz.u8.f16 %0, %1;" : "=r"(tmp) : "r"(__HALF_TO_CUS(h))); + i = static_cast(tmp); +#else + const float f = __half2float(h); + const unsigned char max_val = 0xffU; + const unsigned char min_val = 0U; + const unsigned short bits = + static_cast(static_cast<__half_raw>(h).x << 1U); + // saturation fixup + if (bits > (unsigned short)0xF800U) { + // NaN + i = 0U; + } else if (f > static_cast(max_val)) { + // saturate maximum + i = max_val; + } else if (f < static_cast(min_val)) { + // saturate minimum + i = min_val; + } else { + // normal value, conversion is well-defined + i = static_cast(f); + } +#endif + return i; +} + +__BANG_HOSTDEVICE_FP16_DECL__ signed char __half2char_rz(const __half h) { + signed char i; +#ifdef __BANG_ARCH__ + unsigned int tmp; + asm("cvt.gpr.tz.s8.f16 %0, %1;" : "=r"(tmp) : "r"(__HALF_TO_CUS(h))); + const unsigned char u = static_cast(tmp); + i = static_cast(u); +#else + const float f = __half2float(h); + const signed char max_val = (signed char)0x7fU; + const signed char min_val = (signed char)0x80U; + const unsigned short bits = + static_cast(static_cast<__half_raw>(h).x << 1U); + // saturation fixup + if (bits > (unsigned short)0xF800U) { + // NaN + i = 0; + } else if (f > static_cast(max_val)) { + // saturate maximum + i = max_val; + } else if (f < static_cast(min_val)) { + // saturate minimum + i = min_val; + } else { + // normal value, conversion is well-defined + i = static_cast(f); + } +#endif + return i; +} + +__BANG_HOSTDEVICE_FP16_DECL__ short int __half2short_rz(const __half h) { + short int i; +#ifdef __BANG_ARCH__ + asm("cvt.gpr.tz.s16.f16 %0, %1;" : "+r"(i) : "r"(__HALF_TO_CUS(h))); +#else + const float f = __half2float(h); + const short int max_val = (short int)0x7fffU; + const short int min_val = (short int)0x8000U; + const unsigned short bits = + static_cast(static_cast<__half_raw>(h).x << 1U); + // saturation fixup + if (bits > (unsigned short)0xF800U) { + // NaN + i = 0; + } else if (f > static_cast(max_val)) { + // saturate maximum + i = max_val; + } else if (f < static_cast(min_val)) { + // saturate minimum + i = min_val; + } else { + // normal value, conversion is well-defined + i = static_cast(f); + } +#endif + return i; +} + +/****************************************************************************** +* __half arithmetic * +******************************************************************************/ +#define ___BANG_FP16_STRINGIFY_INNERMOST(x) #x +#define __BANG_FP16_STRINGIFY(x) ___BANG_FP16_STRINGIFY_INNERMOST(x) + +/* Macros for half binary arithmetic */ +#define __BINARY_OP_HALF_MACRO(name) /* do */ {\ + __half val; \ + asm(__BANG_FP16_STRINGIFY(name) ".gpr.f16 %0, %1, %2;" \ + :"+r"(__HALF_TO_US(val)) : "r"(__HALF_TO_CUS(a)),"r"(__HALF_TO_CUS(b))); \ + return val; \ +} /* while(0) */ + +__BANG_HOSTDEVICE_FP16_DECL__ __half __hadd(const __half a, const __half b) { +#ifdef __BANG_ARCH__ + __BINARY_OP_HALF_MACRO(add); +#else + const float fa = __half2float(a); + const float fb = __half2float(b); + return __float2half__(fa + fb); +#endif +} + +__BANG_HOSTDEVICE_FP16_DECL__ __half __hsub(const __half a, const __half b) { +#ifdef __BANG_ARCH__ + __BINARY_OP_HALF_MACRO(sub); +#else + const float fa = __half2float(a); + const float fb = __half2float(b); + return __float2half__(fa - fb); +#endif +} + +__BANG_HOSTDEVICE_FP16_DECL__ __half __hmul(const __half a, const __half b) { +#ifdef __BANG_ARCH__ + __BINARY_OP_HALF_MACRO(mul); +#else + const float fa = __half2float(a); + const float fb = __half2float(b); + return __float2half__(fa * fb); +#endif +} + +__BANG_HOSTDEVICE_FP16_DECL__ __half __hdiv(const __half a, const __half b) { +#ifdef __BANG_ARCH__ + __BINARY_OP_HALF_MACRO(div); +#else + const float fa = __half2float(a); + const float fb = __half2float(b); + return __float2half__(fa / fb); +#endif +} + +__BANG_HOSTDEVICE_FP16_DECL__ __half __hneg(const __half a) { +#ifdef __BANG_ARCH__ + __half r; + asm("neg.gpr.f16 %0, %1;" :"+r"(__HALF_TO_US(r)) : "r"(__HALF_TO_CUS(a))); + return r; +#else + const float fa = __half2float(a); + return __float2half__(-fa); +#endif +} + +/****************************************************************************** +* __half comparison * +******************************************************************************/ +#define __COMPARISON_OP_HALF_MACRO(name, a, b) /* do */ {\ + char val; \ + asm( "{ .gpr %%__$temp3;\n" \ + " " __BANG_FP16_STRINGIFY(name) ".gpr.f16 %%__$temp3, %1, %2;\n" \ + " cvt.gpr.tz.s8.f16 %0, %%__$temp3;}" \ + : "+r"(val) : "r"(__HALF_TO_CUS(a)), "r"(__HALF_TO_CUS(b))); \ + return val; \ +} /* while(0) */ + +__BANG_HOSTDEVICE_FP16_DECL__ bool __heq(const __half a, const __half b) { +#ifdef __BANG_ARCH__ + __COMPARISON_OP_HALF_MACRO(eq, a, b); +#else + const float fa = __half2float(a); + const float fb = __half2float(b); + return (fa == fb); +#endif +} + +__BANG_HOSTDEVICE_FP16_DECL__ bool __hne(const __half a, const __half b) { +#ifdef __BANG_ARCH__ + __COMPARISON_OP_HALF_MACRO(ne, a, b); +#else + const float fa = __half2float(a); + const float fb = __half2float(b); + return (fa != fb) && (!__hisnan(a)) && (!__hisnan(b)); +#endif +} + +__BANG_HOSTDEVICE_FP16_DECL__ bool __hle(const __half a, const __half b) { +#ifdef __BANG_ARCH__ + __COMPARISON_OP_HALF_MACRO(ge, b, a); +#else + const float fa = __half2float(a); + const float fb = __half2float(b); + return (fa <= fb); +#endif +} + +__BANG_HOSTDEVICE_FP16_DECL__ bool __hge(const __half a, const __half b) { +#ifdef __BANG_ARCH__ + __COMPARISON_OP_HALF_MACRO(ge, a, b); +#else + const float fa = __half2float(a); + const float fb = __half2float(b); + return (fa >= fb); +#endif +} + +__BANG_HOSTDEVICE_FP16_DECL__ bool __hlt(const __half a, const __half b) { +#ifdef __BANG_ARCH__ + __COMPARISON_OP_HALF_MACRO(lt, a, b); +#else + const float fa = __half2float(a); + const float fb = __half2float(b); + return (fa < fb); +#endif +} + +__BANG_HOSTDEVICE_FP16_DECL__ bool __hgt(const __half a, const __half b) { +#ifdef __BANG_ARCH__ + __COMPARISON_OP_HALF_MACRO(lt, b, a); +#else + const float fa = __half2float(a); + const float fb = __half2float(b); + return (fa > fb); +#endif +} + +__BANG_HOSTDEVICE_FP16_DECL__ bool __hneu(const __half a, const __half b) { +#ifdef __BANG_ARCH__ + __COMPARISON_OP_HALF_MACRO(neu, a, b); +#else + const float fa = __half2float(a); + const float fb = __half2float(b); + return (fa != fb); +#endif +} + +__BANG_HOSTDEVICE_FP16_DECL__ bool __hisnan(const __half a) { +#if (__BANG_ARCH__ >= 600) + bool r; + asm("test.nan.gpr.f16 %0, %1;" : "+r"(r) : "r"(__HALF_TO_CUS(a))); + return r; +#else + const __half_raw hr = static_cast<__half_raw>(a); + return ((hr.x & (unsigned short)0x7FFFU) > (unsigned short)0x7C00U); +#endif +} + +#undef __BINARY_OP_HALF_MACRO +#undef ___BANG_FP16_STRINGIFY_INNERMOST +#undef __BANG_FP16_STRINGIFY +#undef __COMPARISON_OP_HALF_MACRO + +#undef __BANG_HOSTDEVICE_FP16_DECL__ +#undef __BANG_FP16_DECL__ + +#undef __HALF_TO_US +#undef __HALF_TO_CUS +#undef __BANG_FP16_CONSTEXPR__ + +#if defined(__CPP_VERSION_AT_LEAST_11_FP16) +#undef __CPP_VERSION_AT_LEAST_11_FP16 +#endif /* defined(__CPP_VERSION_AT_LEAST_11_FP16) */ + +#endif /* end of include guard: __BANG_FP16_HPP__ */